fix(task): publish the terminal observation, not the episode's first one - #33
Open
geng-haoran wants to merge 1 commit into
Open
fix(task): publish the terminal observation, not the episode's first one#33geng-haoran wants to merge 1 commit into
geng-haoran wants to merge 1 commit into
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
terminatedis all-False, so its torch.where kept the oldvalue 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,obsstill holds thepre-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: D401removed alongside is a drive-by in the same file: D401 is not enabledin 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