Skip to content

Fix PRNG key reuse in DM Control resets - #349

Open
anxkhn wants to merge 1 commit into
google-deepmind:mainfrom
anxkhn:fix/dm-control-reset-prng-reuse
Open

anxkhn wants to merge 1 commit into
google-deepmind:mainfrom
anxkhn:fix/dm-control-reset-prng-reuse

Conversation

@anxkhn

@anxkhn anxkhn commented Aug 13, 2026

Copy link
Copy Markdown

Summary

Use independent JAX PRNG keys for randomized state components in Finger and Cartpole resets.

Finger resets reused one key for joint positions and the spinner angle. Cartpole Balance reused entropy between position and velocity initialization. Each random draw now receives a distinct split key.

This changes concrete initial states for a given seed, but removes unintended correlation from the reset distribution.

Tests

  • Added key-reuse checks for FingerSpin, FingerTurnEasy, and CartpoleBalance
  • Verified resets under jax.debug_key_reuse(True)

@anxkhn

anxkhn commented Aug 13, 2026

Copy link
Copy Markdown
Author

hey @btaba, could you please review this pr when you have a chance?

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