Skip to content

Add a minimal Fast env for unit testing #306

@YannBerthelot

Description

@YannBerthelot

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

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions