Copybara import of the project:

--
584f657 by andrew <[email protected]>:

rscope supporting changes

--
75b9c39 by andrew <[email protected]>:

run formatter

--
a46f5e1 by andrew <[email protected]>:

elaborate on _model_assets

--
d3500c1 by andrew <[email protected]>:

add deterministic eval option for train script

--
1635951 by andrew <[email protected]>:

modify README.md

--
e602864 by andrew <[email protected]>:

update naming convention

--
3bdb45f by andrew <[email protected]>:

run formatting

--
a145b8f by andrew <[email protected]>:

add option to skip eval rollouts

--
b86274d by andrew <[email protected]>:

update readme

--
2ac436f by andrew <[email protected]>:

move trace_fn out of train_jax_ppo.py

--
1fa8c32 by andrew <[email protected]>:

switch to OOP rscope interface

--
47125c7 by andrew <[email protected]>:

switch training script to the new metrics logging api

--
ce537d7 by andrew <[email protected]>:

improve reward logging

--
5deaaa9 by andrew <[email protected]>:

improve training printouts

--
67f93ba by andrew <[email protected]>:

resolve PR comments

--
0114224 by andrew <[email protected]>:

modify README

COPYBARA_INTEGRATE_REVIEW=#129 from Andrew-Luo1:rscope_new 0114224
PiperOrigin-RevId: 764350162
Change-Id: I701bfc55a1ac994d3440d2d2fe03985ac3f399d5

Author: Andrew-Luo1

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/dm_control_suite/__init__.py

@@ -150,6 +150,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/cartpole.py

@@ -90,8 +90,9 @@ def __init__(
       self._get_reward = self._dense_reward
 
     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)

mujoco_playground/_src/locomotion/__init__.py

@@ -182,7 +182,9 @@ 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/locomotion/t1/randomize.py

@@ -18,7 +18,6 @@
 from mujoco import mjx
 import numpy as np
 
-
 FLOOR_GEOM_ID = 0
 TORSO_BODY_ID = 1
 ANKLE_JOINT_IDS = np.array([[21, 22, 27, 28]])
@@ -30,7 +29,7 @@ def rand_dynamics(rng):
     # Floor friction: =U(0.4, 1.0).
     rng, key = jax.random.split(rng)
     geom_friction = model.geom_friction.at[FLOOR_GEOM_ID, 0].set(
-        jax.random.uniform(key, minval=0.2, maxval=.6)
+        jax.random.uniform(key, minval=0.2, maxval=0.6)
     )
 
     rng, key = jax.random.split(rng)

mujoco_playground/_src/manipulation/__init__.py

@@ -111,7 +111,9 @@ 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/manipulation/franka_emika_panda/panda.py

@@ -71,7 +71,8 @@ def __init__(
 
     self._xml_path = xml_path.as_posix()
     xml = xml_path.read_text()
-    mj_model = mujoco.MjModel.from_xml_string(xml, assets=get_assets())
+    self._model_assets = get_assets()
+    mj_model = mujoco.MjModel.from_xml_string(xml, assets=self._model_assets)
     mj_model.opt.timestep = self.sim_dt
 
     self._mj_model = mj_model
@@ -108,7 +109,7 @@ def _post_init(self, obj_name: str, keyframe: str):
 
   @property
   def xml_path(self) -> str:
-    raise self._xml_path
+    return self._xml_path
 
   @property
   def action_size(self) -> int:

mujoco_playground/_src/mjx_env.py

@@ -273,6 +273,16 @@ def observation_size(self) -> ObservationSize:
       return jax.tree_util.tree_map(lambda x: x.shape, obs)
     return obs.shape[-1]
 
+  @property
+  def model_assets(self) -> Dict[str, Any]:
+    """Dictionary of model assets to use with MjModel.from_xml_path"""
+    if hasattr(self, "_model_assets"):
+      return self._model_assets
+    raise NotImplementedError(
+        "_model_assets not defined for this environment"
+        "see cartpole.py for an example."
+    )
+
   def render(
       self,
       trajectory: List[State],