brax icon indicating copy to clipboard operation
brax copied to clipboard

Update DomainRandomizationVmapWrapper (google#609)

Open AndruGomes13 opened this issue 6 months ago • 3 comments

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.

AndruGomes13 avatar May 17 '25 16:05 AndruGomes13