google-deepmind / google-deepmind/mujoco_playground

Add a minimal `Fast` env for unit testing

Open
#306 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
2.2k
Forks
359
Avg merge
7d 3h
Merged PRs (30d)
1

Description

## Motivation

`mujoco_playground` currently has no toy/stub environment analogous to brax's [`Fast`](https://github.com/google/brax/blob/main/brax/envs/fast.py). 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`](https://github.com/google/brax/blob/main/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

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.