Motivation
mujoco_playground currently has no toy/stub environment analogous to brax's Fast. As a result, any downstream library that wants to cover the playground code path in its unit tests must load a real mjx env (e.g. CartpoleBalance), which incurs:
- mjx / warp kernel compilation on first use,
- non-trivial per-step cost even at
episode_length=5,
- a hard dependency on a physics backend for tests that only care about the wrapper / integration code path.
For integrators (Ajax, TargetGym, SafetyExperiments, etc.) this means either (a) paying the compile cost in CI for every test run, or (b) skipping playground coverage entirely in fast unit tests. Brax's Fast env exists precisely to avoid this tradeoff, it has trivial dynamics, minimal obs/action dims, and no physics backend.
Proposal
Add a Fast-equivalent env to the playground registry with:
- Small fixed obs (e.g. shape
(2,)) and action (e.g. shape (1,)).
- Trivial analytic dynamics, no
mjx.step, no warp kernel compilation.
- Matching
_get_obs(data, info) signature so it behaves like the rest of the registry from a wrapper's perspective.
- A stable id (e.g.
"Fast" or "FastBalance") registered via the normal registry.load(...) path.
Prior art
- brax:
brax/envs/fast.py, used throughout brax's own test suite as the default env for wrapper / training loop correctness tests.
Impact
- Downstream libraries gain a cheap playground path for unit tests (no mjx / warp compilation).
- Playground's own wrapper tests could also benefit, reducing CI time.
- No impact on real-env users, purely additive.
Notes for implementers
- Keep the
_get_obs(data, info) two-arg signature so wrappers written against the general playground convention work unchanged.
- Ensure
BraxAutoResetWrapper + EpisodeWrapper + VmapWrapper compose cleanly around it (i.e. state.info['truncation'] and state.done behave correctly at episode_length boundaries).
I can handle the PR if you feel like this could be a useful addition.
Yann
Motivation
mujoco_playgroundcurrently has no toy/stub environment analogous to brax'sFast. As a result, any downstream library that wants to cover the playground code path in its unit tests must load a real mjx env (e.g.CartpoleBalance), which incurs:episode_length=5,For integrators (Ajax, TargetGym, SafetyExperiments, etc.) this means either (a) paying the compile cost in CI for every test run, or (b) skipping playground coverage entirely in fast unit tests. Brax's
Fastenv exists precisely to avoid this tradeoff, it has trivial dynamics, minimal obs/action dims, and no physics backend.Proposal
Add a
Fast-equivalent env to the playground registry with:(2,)) and action (e.g. shape(1,)).mjx.step, no warp kernel compilation._get_obs(data, info)signature so it behaves like the rest of the registry from a wrapper's perspective."Fast"or"FastBalance") registered via the normalregistry.load(...)path.Prior art
brax/envs/fast.py, used throughout brax's own test suite as the default env for wrapper / training loop correctness tests.Impact
Notes for implementers
_get_obs(data, info)two-arg signature so wrappers written against the general playground convention work unchanged.BraxAutoResetWrapper+EpisodeWrapper+VmapWrappercompose cleanly around it (i.e.state.info['truncation']andstate.donebehave correctly atepisode_lengthboundaries).I can handle the PR if you feel like this could be a useful addition.
Yann