Skip to content

Return current observations from the RSL-RL wrapper - #354

Open
sylvesterkaczmarek wants to merge 1 commit into
google-deepmind:mainfrom
sylvesterkaczmarek:fix/rsl-current-observations
Open

sylvesterkaczmarek wants to merge 1 commit into
google-deepmind:mainfrom
sylvesterkaczmarek:fix/rsl-current-observations

Conversation

@sylvesterkaczmarek

Copy link
Copy Markdown

Summary

RSL-RL's VecEnv.get_observations() contract is to return the current observations. RSLRLBraxWrapper.get_observations() currently delegates directly to reset(), so querying observations resets every wrapped environment.

This matters because RSL-RL queries observations when constructing OnPolicyRunner and again at the start of learn(). With the current wrapper, the second query resets the environments again immediately before rollout.

This change:

  • extracts the existing observation conversion into _current_observations();
  • preserves initialization by calling reset() when env_state is still None;
  • returns observations from the existing env_state on subsequent get_observations() calls;
  • keeps explicit reset() behavior unchanged.

RSL-RL references:

Validation

  • Added a regression test verifying that get_observations() does not call reset() once the environment is initialized.
  • Added a complementary test verifying that an uninitialized wrapper still initializes through reset().
  • Audited the final branch diff: 2 files, +61/-5 lines, with no unrelated changes.
  • Squashed to one focused commit: 522453f80c66e8099419fcd1c3e85d4778fd8496.
  • Branch is based on current upstream main at e74217bb89c77a74ba02e4789263991864375799.
  • Full tests were not run locally in this environment. The focused regression tests are included for upstream CI validation.

This branch has not been deployed

No deployments
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