Skip to content

fix(task): publish the terminal observation, not the episode's first one - #33

Open
geng-haoran wants to merge 1 commit into
mainfrom
fix/terminal-observation
Open

fix(task): publish the terminal observation, not the episode's first one#33
geng-haoran wants to merge 1 commit into
mainfrom
fix/terminal-observation

Conversation

@geng-haoran

Copy link
Copy Markdown
Contributor

RLTaskEnv.step auto-resets done envs in place, so the obs it returns already holds the next
episode's first observation for those envs. info["observations"]["raw"]["obs"] exists to give
off-policy learners the observation the episode actually ended in, which is what a truncated
episode must bootstrap from: V(s_T) for a time-out is a real value, V(reset state) is not.

It did not do that. _raw_observation_cache was written in exactly two places -- reset(), and the
done branch, which stored the post-reset obs. The remaining update lived in an else branch that
runs only when no env is done, where terminated is all-False, so its torch.where kept the old
value unconditionally. Nothing ever advanced the cache during an episode. The raw key therefore
carried the episode's first observation, from reset until the next reset -- not the terminal
obs, and not even the post-reset obs. It was not a stale value; it was a constant.

Verified against the real step() with a counting stub: at a time-out on step 1, the true terminal
obs is 1.0 and the key published 0.0.

The cache never needed to exist. At the point step() builds info, obs still holds the
pre-reset observation for every env -- the auto-reset below is what clobbers it. Snapshotting it
there is both correct and simpler, so the cache is removed rather than repaired.

Downstream, RoboVerse's fast_td3, clean_rl/td3 and clean_rl/sac all read this key as the "true
next obs" for truncation bootstrapping, and were reading a constant for every plain RLTaskEnv
task. Locomotion has not obviously suffered only because LeggedRobotTask overrides step() and
sets the key itself.

The unused # noqa: D401 removed alongside is a drive-by in the same file: D401 is not enabled
in this repo's ruff config, so it fails RUF100 under the pinned ruff 0.14.5 today.

Review: independently reviewed against current main (verdict MERGE AS-IS); rebased, re-verified: 6 passed in 1.02s

🤖 Generated with Claude Code

https://claude.ai/code/session_017i6VtKoovBNed815mWFqxw

RLTaskEnv.step auto-resets done envs in place, so the obs it returns already holds the next
episode's first observation for those envs. info["observations"]["raw"]["obs"] exists to give
off-policy learners the observation the episode actually *ended* in, which is what a truncated
episode must bootstrap from: V(s_T) for a time-out is a real value, V(reset state) is not.

It did not do that. _raw_observation_cache was written in exactly two places -- reset(), and the
done branch, which stored the *post-reset* obs. The remaining update lived in an else branch that
runs only when no env is done, where `terminated` is all-False, so its torch.where kept the old
value unconditionally. Nothing ever advanced the cache during an episode. The raw key therefore
carried the episode's **first** observation, from reset until the next reset -- not the terminal
obs, and not even the post-reset obs. It was not a stale value; it was a constant.

Verified against the real step() with a counting stub: at a time-out on step 1, the true terminal
obs is 1.0 and the key published 0.0.

The cache never needed to exist. At the point step() builds `info`, `obs` still holds the
pre-reset observation for every env -- the auto-reset below is what clobbers it. Snapshotting it
there is both correct and simpler, so the cache is removed rather than repaired.

Downstream, RoboVerse's fast_td3, clean_rl/td3 and clean_rl/sac all read this key as the "true
next obs" for truncation bootstrapping, and were reading a constant for every plain RLTaskEnv
task. Locomotion has not obviously suffered only because LeggedRobotTask overrides step() and
sets the key itself.

The unused `# noqa: D401` removed alongside is a drive-by in the same file: D401 is not enabled
in this repo's ruff config, so it fails RUF100 under the pinned ruff 0.14.5 today.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant