Merge branch 'google-deepmind:main' into add-quadruped

Author: yardenas

.github/workflows/pypi.yml

@@ -0,0 +1,28 @@
+name: Upload Python Package
+
+on:
+  release:
+    types: [created]
+
+jobs:
+  deploy:
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v2
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.10"
+      - name: Install dependencies
+        run: |
+          pip install uv
+          uv pip install --system -e ".[dev]"
+          uv pip install --system build twine
+      - name: Build and publish
+        env:
+          TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }}
+          TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }}
+        run: |
+          python -m build
+          twine upload --username $TWINE_USERNAME --password $TWINE_PASSWORD dist/*

CHANGELOG.md

@@ -2,6 +2,24 @@
 
 All notable changes to this project will be documented in this file.
 
+## Next release
+
+- Pass through the [MuJoCo Warp](https://github.com/google-deepmind/mujoco_warp)
+  (MjWarp) implementation to MJX, so that MuJoCo Playground environments can
+  train with MuJoCo Warp! DM Control Suite and Locomotion environments now
+  support MjWarp. You can pass through the implementation via the config
+  override
+  `registry.load('CartpoleBalance', config_overrides={'impl': 'warp'})`.
+
+## [0.0.5] - 2025-06-23
+
+- Change `light_directional` to `light_type` following MuJoCo API change from version 3.3.2 to 3.3.3. Fixes https://github.com/google-deepmind/mujoco_playground/issues/142.
+- Fix bug in `get_qpos_ids`.
+- Implement `render` in Wrapper.
+- Fix https://github.com/google-deepmind/mujoco_playground/issues/123.
+- Fix https://github.com/google-deepmind/mujoco_playground/issues/126.
+- Fix https://github.com/google-deepmind/mujoco_playground/issues/41.
+
 ## [0.0.4] - 2025-02-07
 
 ### Added

README.md

@@ -63,6 +63,24 @@ For vision-based environments, please refer to the installation instructions in
 | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco_playground/blob/main/learning/notebooks/training_vision_1.ipynb) | Training CartPole from Vision |
 | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco_playground/blob/main/learning/notebooks/training_vision_2.ipynb) | Robotic Manipulation from Vision |
 
+## Running from CLI
+> [!IMPORTANT]
+> Assumes installation from source.
+
+For basic usage, navigate to the repo's directory and run:
+```bash
+python learning/train_jax_ppo.py --env_name CartpoleBalance
+```
+
+### Training Visualization
+
+To interactively view trajectories throughout training with [rscope](https://github.com/Andrew-Luo1/rscope/tree/main), install it (`pip install rscope`) and run:
+
+```
+python learning/train_jax_ppo.py --env_name PandaPickCube --rscope_envs 16 --run_evals=False --deterministic_rscope=True
+# In a separate terminal
+python -m rscope
+```
 
 ## FAQ
 

learning/train_jax_ppo.py

@@ -132,6 +132,33 @@
     "policy_obs_key", "state", "Policy obs key"
 )
 _VALUE_OBS_KEY = flags.DEFINE_string("value_obs_key", "state", "Value obs key")
+_RSCOPE_ENVS = flags.DEFINE_integer(
+    "rscope_envs",
+    None,
+    "Number of parallel environment rollouts to save for the rscope viewer",
+)
+_DETERMINISTIC_RSCOPE = flags.DEFINE_boolean(
+    "deterministic_rscope",
+    True,
+    "Run deterministic rollouts for the rscope viewer",
+)
+_RUN_EVALS = flags.DEFINE_boolean(
+    "run_evals",
+    True,
+    "Run evaluation rollouts between policy updates.",
+)
+_LOG_TRAINING_METRICS = flags.DEFINE_boolean(
+    "log_training_metrics",
+    False,
+    "Whether to log training metrics and callback to progress_fn. Significantly"
+    " slows down training if too frequent.",
+)
+_TRAINING_METRICS_STEPS = flags.DEFINE_integer(
+    "training_metrics_steps",
+    1_000_000,
+    "Number of steps between logging training metrics. Increase if training"
+    " experiences slowdown.",
+)
 
 
 def get_rl_config(env_name: str) -> config_dict.ConfigDict:
@@ -151,6 +178,24 @@ def get_rl_config(env_name: str) -> config_dict.ConfigDict:
   raise ValueError(f"Env {env_name} not found in {registry.ALL_ENVS}.")
 
 
+def rscope_fn(full_states, obs, rew, done):
+  """
+  All arrays are of shape (unroll_length, rscope_envs, ...)
+  full_states: dict with keys 'qpos', 'qvel', 'time', 'metrics'
+  obs: nd.array or dict obs based on env configuration
+  rew: nd.array rewards
+  done: nd.array done flags
+  """
+  # Calculate cumulative rewards per episode, stopping at first done flag
+  done_mask = jp.cumsum(done, axis=0)
+  valid_rewards = rew * (done_mask == 0)
+  episode_rewards = jp.sum(valid_rewards, axis=0)
+  print(
+      "Collected rscope rollouts with reward"
+      f" {episode_rewards.mean():.3f} +- {episode_rewards.std():.3f}"
+  )
+
+
 def main(argv):
   """Run training and evaluation for the specified environment."""
 
@@ -209,11 +254,16 @@ def main(argv):
     ppo_params.network_factory.policy_obs_key = _POLICY_OBS_KEY.value
   if _VALUE_OBS_KEY.present:
     ppo_params.network_factory.value_obs_key = _VALUE_OBS_KEY.value
-
   if _VISION.value:
     env_cfg.vision = True
     env_cfg.vision_config.render_batch_size = ppo_params.num_envs
   env = registry.load(_ENV_NAME.value, config=env_cfg)
+  if _RUN_EVALS.present:
+    ppo_params.run_evals = _RUN_EVALS.value
+  if _LOG_TRAINING_METRICS.present:
+    ppo_params.log_training_metrics = _LOG_TRAINING_METRICS.value
+  if _TRAINING_METRICS_STEPS.present:
+    ppo_params.training_metrics_steps = _TRAINING_METRICS_STEPS.value
 
   print(f"Environment Config:\n{env_cfg}")
   print(f"PPO Training Parameters:\n{ppo_params}")
@@ -268,13 +318,6 @@ def main(argv):
   with open(ckpt_path / "config.json", "w", encoding="utf-8") as fp:
     json.dump(env_cfg.to_dict(), fp, indent=4)
 
-  # Define policy parameters function for saving checkpoints
-  def policy_params_fn(current_step, make_policy, params):  # pylint: disable=unused-argument
-    orbax_checkpointer = ocp.PyTreeCheckpointer()
-    save_args = orbax_utils.save_args_from_target(params)
-    path = ckpt_path / f"{current_step}"
-    orbax_checkpointer.save(path, params, force=True, save_args=save_args)
-
   training_params = dict(ppo_params)
   if "network_factory" in training_params:
     del training_params["network_factory"]
@@ -319,9 +362,9 @@ def policy_params_fn(current_step, make_policy, params):  # pylint: disable=unus
       ppo.train,
       **training_params,
       network_factory=network_factory,
-      policy_params_fn=policy_params_fn,
       seed=_SEED.value,
       restore_checkpoint_path=restore_checkpoint_path,
+      save_checkpoint_path=ckpt_path,
       wrap_env_fn=None if _VISION.value else wrapper.wrap_for_brax_training,
       num_eval_envs=num_eval_envs,
   )
@@ -341,18 +384,55 @@ def progress(num_steps, metrics):
       for key, value in metrics.items():
         writer.add_scalar(key, value, num_steps)
       writer.flush()
-
-    print(f"{num_steps}: reward={metrics['eval/episode_reward']:.3f}")
+    if _RUN_EVALS.value:
+      print(f"{num_steps}: reward={metrics['eval/episode_reward']:.3f}")
+    if _LOG_TRAINING_METRICS.value:
+      if "episode/sum_reward" in metrics:
+        print(
+            f"{num_steps}: mean episode"
+            f" reward={metrics['episode/sum_reward']:.3f}"
+        )
 
   # Load evaluation environment
   eval_env = (
       None if _VISION.value else registry.load(_ENV_NAME.value, config=env_cfg)
   )
 
+  policy_params_fn = lambda *args: None
+  if _RSCOPE_ENVS.value:
+    # Interactive visualisation of policy checkpoints
+    from rscope import brax as rscope_utils
+
+    if not _VISION.value:
+      rscope_env = registry.load(_ENV_NAME.value, config=env_cfg)
+      rscope_env = wrapper.wrap_for_brax_training(
+          rscope_env,
+          episode_length=ppo_params.episode_length,
+          action_repeat=ppo_params.action_repeat,
+          randomization_fn=training_params.get("randomization_fn"),
+      )
+    else:
+      rscope_env = env
+
+    rscope_handle = rscope_utils.BraxRolloutSaver(
+        rscope_env,
+        ppo_params,
+        _VISION.value,
+        _RSCOPE_ENVS.value,
+        _DETERMINISTIC_RSCOPE.value,
+        jax.random.PRNGKey(_SEED.value),
+        rscope_fn,
+    )
+
+    def policy_params_fn(current_step, make_policy, params):  # pylint: disable=unused-argument
+      rscope_handle.set_make_policy(make_policy)
+      rscope_handle.dump_rollout(params)
+
   # Train or load the model
   make_inference_fn, params, _ = train_fn(  # pylint: disable=no-value-for-parameter
       environment=env,
       progress_fn=progress,
+      policy_params_fn=policy_params_fn,
       eval_env=None if _VISION.value else eval_env,
   )
 

mujoco_playground/__init__.py

@@ -25,6 +25,7 @@
 from mujoco_playground._src.mjx_env import render_array
 from mujoco_playground._src.mjx_env import State
 from mujoco_playground._src.mjx_env import step
+
 # pylint: enable=g-importing-member
 
 __all__ = [

mujoco_playground/_src/collision.py

@@ -17,22 +17,29 @@
 from typing import Any, Tuple
 
 import jax
-import jax.numpy as jnp
+import jax.numpy as jp
 from mujoco import mjx
+from mujoco.mjx._src import types
 
 
 def get_collision_info(
     contact: Any, geom1: int, geom2: int
 ) -> Tuple[jax.Array, jax.Array]:
   """Get the distance and normal of the collision between two geoms."""
-  mask = (jnp.array([geom1, geom2]) == contact.geom).all(axis=1)
-  mask |= (jnp.array([geom2, geom1]) == contact.geom).all(axis=1)
-  idx = jnp.where(mask, contact.dist, 1e4).argmin()
+  mask = (jp.array([geom1, geom2]) == contact.geom).all(axis=1)
+  mask |= (jp.array([geom2, geom1]) == contact.geom).all(axis=1)
+  idx = jp.where(mask, contact.dist, 1e4).argmin()
   dist = contact.dist[idx] * mask[idx]
   normal = (dist < 0) * contact.frame[idx, 0, :3]
   return dist, normal
 
 
 def geoms_colliding(state: mjx.Data, geom1: int, geom2: int) -> jax.Array:
   """Return True if the two geoms are colliding."""
-  return get_collision_info(state.contact, geom1, geom2)[0] < 0
+  # if not isinstance(state._impl, types.DataJAX):
+  #   raise NotImplementedError(
+  #       "`geoms_colliding` only implemented for JAX MJX backend."
+  #   )
+  if not isinstance(state._impl, types.DataJAX):
+    return jp.array(False)
+  return get_collision_info(state._impl.contact, geom1, geom2)[0] < 0  # pylint: disable=protected-access

mujoco_playground/_src/dm_control_suite/__init__.py

@@ -155,6 +155,8 @@ def load(
       An instance of the environment.
   """
   if env_name not in _envs:
-    raise ValueError(f"Env '{env_name}' not found. Available envs: {_cfgs.keys()}")
+    raise ValueError(
+        f"Env '{env_name}' not found. Available envs: {_cfgs.keys()}"
+    )
   config = config or get_default_config(env_name)
   return _envs[env_name](config=config, config_overrides=config_overrides)

mujoco_playground/_src/dm_control_suite/acrobot.py

@@ -36,6 +36,9 @@ def default_config() -> config_dict.ConfigDict:
       episode_length=1000,
       action_repeat=1,
       vision=False,
+      impl="jax",
+      nconmax=0,
+      njmax=0,
   )
 
 
@@ -57,11 +60,12 @@ def __init__(
     self._margin = 0.0 if sparse else 1.0
 
     self._xml_path = _XML_PATH.as_posix()
+    self._model_assets = common.get_assets()
     self._mj_model = mujoco.MjModel.from_xml_string(
-        _XML_PATH.read_text(), common.get_assets()
+        _XML_PATH.read_text(), self._model_assets
     )
     self._mj_model.opt.timestep = self.sim_dt
-    self._mjx_model = mjx.put_model(self._mj_model)
+    self._mjx_model = mjx.put_model(self._mj_model, impl=self._config.impl)
     self._post_init()
 
   def _post_init(self) -> None:
@@ -77,7 +81,14 @@ def reset(self, rng: jax.Array) -> mjx_env.State:
     qpos = jax.random.uniform(
         rng1, (self.mjx_model.nq,), minval=-jp.pi, maxval=jp.pi
     )
-    data = mjx_env.init(self.mjx_model, qpos=qpos)
+    data = mjx_env.make_data(
+        self.mj_model,
+        qpos=qpos,
+        impl=self.mjx_model.impl.value,
+        nconmax=self._config.nconmax,
+        njmax=self._config.njmax,
+    )
+    data = mjx.forward(self.mjx_model, data)
 
     metrics = {
         "distance": jp.zeros(()),

mujoco_playground/_src/dm_control_suite/ball_in_cup.py

@@ -35,6 +35,9 @@ def default_config() -> config_dict.ConfigDict:
       episode_length=1000,
       action_repeat=1,
       vision=False,
+      impl="jax",
+      nconmax=10_000,
+      njmax=25,
   )
 
 
@@ -53,11 +56,12 @@ def __init__(
       )
 
     self._xml_path = _XML_PATH.as_posix()
+    self._model_assets = common.get_assets()
     self._mj_model = mujoco.MjModel.from_xml_string(
-        _XML_PATH.read_text(), common.get_assets()
+        _XML_PATH.read_text(), self._model_assets
     )
     self._mj_model.opt.timestep = self.sim_dt
-    self._mjx_model = mjx.put_model(self._mj_model)
+    self._mjx_model = mjx.put_model(self._mj_model, impl=self._config.impl)
     self._post_init()
 
   def _post_init(self) -> None:
@@ -68,7 +72,13 @@ def _post_init(self) -> None:
     self._ball_size = self._mj_model.geom_size[geom_id, 0]
 
   def reset(self, rng: jax.Array) -> mjx_env.State:
-    data = mjx_env.init(self.mjx_model)
+    data = mjx_env.make_data(
+        self.mj_model,
+        impl=self.mjx_model.impl.value,
+        nconmax=self._config.nconmax,
+        njmax=self._config.njmax,
+    )
+    data = mjx.forward(self.mjx_model, data)
 
     metrics = {}
     info = {"rng": rng}