feat: pass randomness keys to environment steps

[quinnarnold]

Summary

Allow callers to resample a stochastic successor with env.step(state, action, key), while preserving the existing behavior of env.step(state, action).

For stochastic environments, an explicit key replaces state.key using the environment’s existing splitting and returned-state key behavior. Deterministic environments accept and may ignore the optional key. Wrappers forward it; AutoResetWrapper continues to derive reset randomness from the returned state key.

State.key remains in place. Existing training callers, examples, and Gym/DMEnv adapters retain their calling conventions. Merged current main and resolved the shared helper/fake conflicts while preserving upstream changes.

Closes #211.

Validation

• 124 focused environment, wrapper, and testing-helper tests passed.
• 55 existing environment smoke tests passed.
• After the final wrapper/test simplification, all 63 affected tests passed again.
• Direct comparisons against current main passed for all seven stochastic environments: default calls match upstream, supplied keys act independently of the old stored key, and switching explicit/default calls through JIT works. Wrapper, batched, and terminal automatic-reset checks passed.
• All project pre-commit hooks and whitespace checks passed.
• Retained the PR’s single explicit-key Game2048 regression, trimmed its redundant repeat assertion, and verified it rejects ignored-key and constant-key implementations. No additional test functions added in this revision.

AI-assisted implementation and validation.

[sash-a]

This is such a huge break in the API that I feel we'd need more careful planning and an agreed upon approach to merge this.

key's aren't truly optional here, e.g in 2048 if the key is none you raise an error. For this to be merged we'd really want a way to keep the existing functionality (using the state's key) as the default and allow passing the key as an option.

This is quite tricky in JAX. If we do a cond it's more computational work and if we have an if this will silently fail if you use the same jitted function when passing a key vs using the state's key.

[quinnarnold]

Thanks, I agree the compatibility contract needs clarification before I revise this further. The current implementation breaks existing callers despite the optional signature, including in AutoResetWrapper.

I’d propose keeping step(state, action) unchanged, including its state-key splitting and reset behavior, and making external randomness an opt-in through step(state, action, key). State.key would remain; removing it or migrating existing training callers would be separate work. This would address successor resampling from #211 without requiring the broader API transition taken by PGX.

On the JIT concern, I checked a minimal function using state_key if key is None else key: alternating None and an array through the same jitted function produced two correct traces on JAX 0.10.1, without a runtime cond. That only checks the argument-selection behavior, not the full environment/wrapper path. Is there a particular calling pattern or supported JAX version behind the silent-failure concern that I should examine?

Would this backward-compatible scope be acceptable? The remaining design question is how an explicit key determines the returned state key and any automatic reset, so switching between the two calling forms has defined behavior. I’ll hold further edits pending your feedback on that direction.

[sash-a]

Ok I just tested that if and it does seem to work, I thought it wouldn't. Then I am happy to go through with this using state.key if key isn't passed and using the key argument if it is