brax
brax copied to clipboard
Update DomainRandomizationVmapWrapper (google#609)
This change ensures that the original system (env.sys) remains unchanged during JAX tracing. This prevents the UnexpectedTracerError that would occur when wrapping the same environment twice.
Tests:
- Training on a local personal env works.
- Google Colab tutorial runs without needing the duplicate eval_env workaround.