[Feature] mujoco_playground wrapper review fixes

- CI: install `playground` (PyPI name) instead of import name
  `mujoco_playground`; the previous `pip install mujoco_playground` was the
  root cause of the current `unittests-mujoco-playground` CI failure.
- CI: drop the broken try/except JAX-init fallback in run_test.sh — it ran
  in a subprocess that exited immediately, so `JAX_PLATFORM_NAME` never
  reached pytest and any real "GPU not visible" error was hidden.
- Wrapper: docstrings no longer advertise a `state` field that the env
  never emits; the JAX state is intentionally kept on
  `self._current_state` rather than round-tripped through TensorDict (this
  is now documented as a `.. note::` block).
- Wrapper: freeze `MujocoPlaygroundAgentSpec` /
  `MujocoPlaygroundAgentMapping` dataclasses and deep-copy
  `KNOWN_MARL_MAPPINGS` entries on string lookup so users cannot mutate
  the module-level mapping by accident.
- Wrapper: emit a `UserWarning` when resolving a string against
  `KNOWN_MARL_MAPPINGS`, since those indices target Brax's observation
  layout and may not be semantically equivalent for mujoco_playground envs.
- Wrapper: document the policy contract for `homogenization_mode='max'`
  and `'concat'` (which action/obs entries are real vs padding/discarded).
- Wrapper: align `_MujocoPlaygroundMeta` num_workers handling with
  `_BraxMeta`; make `agent_mapping` and `config`/`config_overrides`
  keyword-only; accept `seed=None` in `_set_seed` (defaults to 0, matching
  the `_reset` fallback) instead of raising bare `Exception`; document
  `_listerize`'s inclusive-range semantics; drop unused
  `pixels_only`/`camera_id`/`render_kwargs` from `_build_env`.
- Example: rewrite `save_visualization` in
  `profile_mujoco_playground_collector.py` to snapshot
  `env._current_state` during a manual rollout, removing dependencies on
  non-existent `env._state_example` and `td["state"]`.
- Config: drop stale `categorical_action_encoding` field on
  `MujocoPlaygroundEnvConfig`; add `agent_mapping` and `num_workers` so
  the MARL and parallel-process knobs are reachable from the config
  system.
- Tests: replace three duplicated `_setup_jax` fixtures with a single
  module-level autouse fixture; introduce a session-scoped
  `marl_env_sizes` fixture to avoid constructing a throwaway env in every
  MARL test; expect the new MABrax warning on string lookups; add a
  new `TestMujocoPlaygroundDictObs` class covering reset /
  `check_env_specs` and a negative test that
  `homogenization_mode != 'none'` raises `NotImplementedError` for
  dict-obs envs; make `test_no_mapping_regression` actually exercise the
  MARL env it constructs.
- Docs: keep the `MOGymEnv` / `MOGymWrapper` pair adjacent in
  `envs_libraries.rst`; `__repr__` now uses `env_name=` to match the
  constructor kwarg.

Author: vmoens

.github/unittest/linux_libs/scripts_mujoco_playground/environment.yml

@@ -21,5 +21,6 @@ dependencies:
     - scipy
     - hydra-core
     - jax[cuda12]>=0.7.0
-    - mujoco_playground
+    # NOTE: import name is `mujoco_playground` but the PyPI distribution is `playground`.
+    - playground
     - psutil

.github/unittest/linux_libs/scripts_mujoco_playground/run_test.sh

@@ -30,31 +30,11 @@ export MAGNUM_LOG=verbose MAGNUM_GPU_VALIDATION=ON
 python -c "import mujoco_playground"
 python -c "from mujoco_playground import dm_control_suite, locomotion, manipulation, registry"
 
-# Initialize JAX with proper GPU configuration
-python -c "
-import jax
-import jax.numpy as jnp
-import os
-
-# Configure JAX for GPU
-os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = 'false'
-os.environ['XLA_PYTHON_CLIENT_ALLOCATOR'] = 'platform'
-
-# Test JAX GPU availability
-try:
-    devices = jax.devices()
-    print(f'JAX devices: {devices}')
-    if len(devices) > 1:
-        print('JAX GPU is available')
-    else:
-        print('JAX CPU only')
-except Exception as e:
-    print(f'JAX initialization error: {e}')
-    # Fallback to CPU
-    os.environ['JAX_PLATFORM_NAME'] = 'cpu'
-    jax.config.update('jax_platform_name', 'cpu')
-    print('Falling back to JAX CPU')
-"
+# Report JAX devices. We deliberately avoid try/except + JAX_PLATFORM_NAME
+# fallback here: it runs in a subprocess that exits immediately, so it has no
+# effect on the pytest invocation below, and it would hide a real "GPU not
+# visible" failure from CI.
+python -c "import jax; print(f'JAX devices: {jax.devices()}')"
 
 python -c 'import torch;t = torch.ones([2,2], device="cuda:0");print(t);print("tensor device:" + str(t.device))'
 

docs/source/reference/envs_libraries.rst

@@ -97,9 +97,9 @@ Available wrappers
     MeltingpotEnv
     MeltingpotWrapper
     MOGymEnv
+    MOGymWrapper
     MujocoPlaygroundEnv
     MujocoPlaygroundWrapper
-    MOGymWrapper
     MultiThreadedEnv
     MultiThreadedEnvWrapper
     OpenMLEnv

examples/collectors/profile_mujoco_playground_collector.py

@@ -41,7 +41,7 @@
 from torchrl._utils import logger as torchrl_logger
 from torchrl.collectors import Collector
 from torchrl.envs import MujocoPlaygroundEnv
-from torchrl.envs.libs.jax_utils import _tensordict_to_object
+from torchrl.envs.utils import step_mdp
 
 os.environ.setdefault("TORCHRL_PROFILING", "1")
 
@@ -262,43 +262,44 @@ def make_policy(
     raise ValueError(f"Unknown policy mode {policy_mode!r}.")
 
 
-def save_visualization(env: MujocoPlaygroundEnv, args: argparse.Namespace) -> None:
+def save_visualization(
+    env: MujocoPlaygroundEnv,
+    policy: typing.Callable[[TensorDictBase], TensorDictBase],
+    args: argparse.Namespace,
+) -> None:
     """Collect a short rollout and save a rendered visualization.
 
-    Converts TensorDict states back to JAX State objects, calls the
-    underlying environment's ``render`` method, then writes the frames
-    to an MP4 file (via ``imageio``) or to individual PNG files.
+    Snapshots the env's JAX state (``env._current_state``) after each step,
+    optionally indexing into the leading batch dim, and feeds the resulting
+    list of states to the underlying env's ``render`` method.
 
     Args:
         env (MujocoPlaygroundEnv): the environment to render.
+        policy (callable): a policy that fills the action in a TensorDict.
         args (argparse.Namespace): parsed command-line arguments.
     """
     import jax
 
+    batched = len(env.batch_size) > 0
+
+    def _snapshot(state):
+        # For batched envs, pick the first sub-environment so the renderer
+        # receives un-batched JAX State objects.
+        if batched:
+            return jax.tree_util.tree_map(lambda x: x[0], state)
+        return state
+
     torchrl_logger.info(
         f"Collecting {args.render_steps} steps for visualization rollout."
     )
+    td = env.reset()
+    states = [_snapshot(env._current_state)]
     with torch.no_grad():
-        td = env.rollout(args.render_steps)
-
-    batched = len(env.batch_size) > 0
-
-    # Build a single-element state example so _tensordict_to_object can
-    # infer field structure without a batch dimension.
-    if batched:
-        single_example = jax.tree_util.tree_map(lambda x: x[0], env._state_example)
-    else:
-        single_example = env._state_example
-
-    torchrl_logger.info("Converting TensorDict states to JAX State objects.")
-    states = []
-    for t in range(args.render_steps):
-        if batched:
-            step_td = td[0, t].get("state")
-        else:
-            step_td = td[t].get("state")
-        jax_state = _tensordict_to_object(step_td, single_example)
-        states.append(jax_state)
+        for _ in range(args.render_steps):
+            td = policy(td)
+            td = env.step(td)
+            states.append(_snapshot(env._current_state))
+            td = step_mdp(td)
 
     torchrl_logger.info(
         f"Rendering {len(states)} frames "
@@ -402,7 +403,7 @@ def main() -> None:
     torchrl_logger.info(f"Trace written to {worker_0_trace_path}.")
 
     if args_cli.render:
-        save_visualization(env, args_cli)
+        save_visualization(env, policy, args_cli)
 
     env.close()