Playground updates for manipulation environments using Warp.

PiperOrigin-RevId: 800484715
Change-Id: I64880e02d73abbf225119e23d862dd162bd41803

Author: btaba

learning/train_jax_ppo.py

@@ -166,16 +166,16 @@
 def get_rl_config(env_name: str) -> config_dict.ConfigDict:
   if env_name in mujoco_playground.manipulation._envs:
     if _VISION.value:
-      return manipulation_params.brax_vision_ppo_config(env_name)
-    return manipulation_params.brax_ppo_config(env_name)
+      return manipulation_params.brax_vision_ppo_config(env_name, _IMPL.value)
+    return manipulation_params.brax_ppo_config(env_name, _IMPL.value)
   elif env_name in mujoco_playground.locomotion._envs:
-    if _VISION.value:
-      return locomotion_params.brax_vision_ppo_config(env_name)
-    return locomotion_params.brax_ppo_config(env_name)
+    return locomotion_params.brax_ppo_config(env_name, _IMPL.value)
   elif env_name in mujoco_playground.dm_control_suite._envs:
     if _VISION.value:
-      return dm_control_suite_params.brax_vision_ppo_config(env_name)
-    return dm_control_suite_params.brax_ppo_config(env_name)
+      return dm_control_suite_params.brax_vision_ppo_config(
+          env_name, _IMPL.value
+      )
+    return dm_control_suite_params.brax_ppo_config(env_name, _IMPL.value)
 
   raise ValueError(f"Env {env_name} not found in {registry.ALL_ENVS}.")
 

mujoco_playground/_src/locomotion/berkeley_humanoid/xmls/scene_mjx_feetonly_rough_terrain.xml

@@ -15,7 +15,7 @@
     <!-- https://polyhaven.com/a/rock_face -->
     <texture type="2d" name="groundplane" file="assets/rocky_texture.png"/>
     <material name="groundplane" texture="groundplane" texuniform="true" texrepeat="5 5" reflectance=".8"/>
-    <hfield name="hfield" file="assets/hfield.png" size="10 10 .05 0.1"/>
+    <hfield name="hfield" file="assets/hfield.png" size="10 10 .05 1.0"/>
   </asset>
 
   <worldbody>

mujoco_playground/_src/locomotion/g1/joystick.py

@@ -103,7 +103,7 @@ def default_config() -> config_dict.ConfigDict:
       ang_vel_yaw=[-1.0, 1.0],
       impl="jax",
       nconmax=8 * 8192,
-      njmax=29 + 8 * 4,
+      njmax=29 * 2 + 8 * 4,
   )
 
 
@@ -118,7 +118,7 @@ def __init__(
   ):
     if task.startswith("rough"):
       config.nconmax = 100 * 8192
-      config.njmax = 29 + 100 * 4
+      config.njmax = 29 * 2 + 100 * 4
     super().__init__(
         xml_path=consts.task_to_xml(task).as_posix(),
         config=config,

mujoco_playground/_src/locomotion/g1/xmls/scene_mjx_feetonly_rough_terrain.xml

@@ -16,7 +16,7 @@
     <!-- https://polyhaven.com/a/rock_face -->
     <texture type="2d" name="groundplane" file="assets/rocky_texture.png"/>
     <material name="groundplane" texture="groundplane" texuniform="true" texrepeat="5 5" reflectance=".8"/>
-    <hfield name="hfield" file="assets/hfield.png" size="10 10 .05 0.1"/>
+    <hfield name="hfield" file="assets/hfield.png" size="10 10 .05 1.0"/>
   </asset>
 
   <worldbody>

mujoco_playground/_src/locomotion/go1/xmls/scene_mjx_feetonly_rough_terrain.xml

@@ -15,7 +15,7 @@
     <!-- https://polyhaven.com/a/rock_face -->
     <texture type="2d" name="groundplane" file="assets/rocky_texture.png"/>
     <material name="groundplane" texture="groundplane" texuniform="true" texrepeat="5 5" reflectance=".8"/>
-    <hfield name="hfield" file="assets/hfield.png" size="10 10 .05 0.1"/>
+    <hfield name="hfield" file="assets/hfield.png" size="10 10 .05 1.0"/>
   </asset>
 
   <worldbody>

mujoco_playground/_src/locomotion/t1/xmls/scene_mjx_feetonly_rough_terrain.xml

@@ -15,7 +15,7 @@
     <!-- https://polyhaven.com/a/rock_face -->
     <texture type="2d" name="groundplane" file="assets/rocky_texture.png"/>
     <material name="groundplane" texture="groundplane" texuniform="true" texrepeat="5 5" reflectance=".8"/>
-    <hfield name="hfield" file="assets/hfield.png" size="10 10 .05 0.1"/>
+    <hfield name="hfield" file="assets/hfield.png" size="10 10 .05 1.0"/>
   </asset>
 
   <worldbody>

mujoco_playground/_src/manipulation/aloha/single_peg_insertion.py

@@ -48,7 +48,7 @@ def default_config() -> config_dict.ConfigDict:
               peg_insertion_reward=8,
           )
       ),
-      impl='jax',
+      impl="jax",
       nconmax=24 * 8192,
       njmax=256,
   )

mujoco_playground/config/dm_control_suite_params.py

@@ -14,12 +14,14 @@
 # ==============================================================================
 """RL config for DM Control Suite."""
 
+from typing import Optional
 from ml_collections import config_dict
-
 from mujoco_playground._src import dm_control_suite
 
 
-def brax_ppo_config(env_name: str) -> config_dict.ConfigDict:
+def brax_ppo_config(
+    env_name: str, impl: Optional[str] = None
+) -> config_dict.ConfigDict:
   """Returns tuned Brax PPO config for the given environment."""
   env_config = dm_control_suite.get_default_config(env_name)
 
@@ -38,6 +40,7 @@ def brax_ppo_config(env_name: str) -> config_dict.ConfigDict:
       entropy_cost=1e-2,
       num_envs=2048,
       batch_size=1024,
+      num_resets_per_eval=10,
   )
 
   if env_name.startswith("AcrobotSwingup"):
@@ -57,7 +60,9 @@ def brax_ppo_config(env_name: str) -> config_dict.ConfigDict:
   return rl_config
 
 
-def brax_vision_ppo_config(env_name: str) -> config_dict.ConfigDict:
+def brax_vision_ppo_config(
+    env_name: str, unused_impl: Optional[str] = None
+) -> config_dict.ConfigDict:
   """Returns tuned Brax Vision PPO config for the given environment."""
   env_config = dm_control_suite.get_default_config(env_name)
 
@@ -80,6 +85,7 @@ def brax_vision_ppo_config(env_name: str) -> config_dict.ConfigDict:
       num_eval_envs=1024,
       batch_size=256,
       max_grad_norm=1.0,
+      num_resets_per_eval=10,
   )
 
   if env_name != "CartpoleBalance":
@@ -88,7 +94,9 @@ def brax_vision_ppo_config(env_name: str) -> config_dict.ConfigDict:
   return rl_config
 
 
-def brax_sac_config(env_name: str) -> config_dict.ConfigDict:
+def brax_sac_config(
+    env_name: str, unused_impl: Optional[str] = None
+) -> config_dict.ConfigDict:
   """Returns tuned Brax SAC config for the given environment."""
   env_config = dm_control_suite.get_default_config(env_name)
 
@@ -109,6 +117,7 @@ def brax_sac_config(env_name: str) -> config_dict.ConfigDict:
       network_factory=config_dict.create(
           q_network_layer_norm=True,
       ),
+      num_resets_per_eval=10,
   )
 
   if env_name == "PendulumSwingUp":

mujoco_playground/config/locomotion_params.py

@@ -14,12 +14,14 @@
 # ==============================================================================
 """RL config for Locomotion envs."""
 
+from typing import Optional
 from ml_collections import config_dict
-
 from mujoco_playground._src import locomotion
 
 
-def brax_ppo_config(env_name: str) -> config_dict.ConfigDict:
+def brax_ppo_config(
+    env_name: str, impl: Optional[str] = None
+) -> config_dict.ConfigDict:
   """Returns tuned Brax PPO config for the given environment."""
   env_config = locomotion.get_default_config(env_name)
 
@@ -45,12 +47,12 @@ def brax_ppo_config(env_name: str) -> config_dict.ConfigDict:
           policy_obs_key="state",
           value_obs_key="state",
       ),
+      num_resets_per_eval=10,
   )
 
   if env_name in ("Go1JoystickFlatTerrain", "Go1JoystickRoughTerrain"):
     rl_config.num_timesteps = 200_000_000
     rl_config.num_evals = 10
-    rl_config.num_resets_per_eval = 1
     rl_config.network_factory = config_dict.create(
         policy_hidden_layer_sizes=(512, 256, 128),
         value_hidden_layer_sizes=(512, 256, 128),
@@ -109,7 +111,6 @@ def brax_ppo_config(env_name: str) -> config_dict.ConfigDict:
     rl_config.num_timesteps = 150_000_000
     rl_config.num_evals = 15
     rl_config.clipping_epsilon = 0.2
-    rl_config.num_resets_per_eval = 1
     rl_config.entropy_cost = 0.005
     rl_config.network_factory = config_dict.create(
         policy_hidden_layer_sizes=(512, 256, 128),
@@ -163,7 +164,9 @@ def brax_ppo_config(env_name: str) -> config_dict.ConfigDict:
   return rl_config
 
 
-def rsl_rl_config(env_name: str) -> config_dict.ConfigDict:
+def rsl_rl_config(
+    env_name: str, unused_impl: Optional[str] = None
+) -> config_dict.ConfigDict:
   """Returns tuned RSL-RL PPO config for the given environment."""
 
   rl_config = config_dict.create(

mujoco_playground/config/manipulation_params.py

@@ -14,12 +14,14 @@
 # ==============================================================================
 """RL config for Manipulation envs."""
 
+from typing import Optional
 from ml_collections import config_dict
-
 from mujoco_playground._src import manipulation
 
 
-def brax_ppo_config(env_name: str) -> config_dict.ConfigDict:
+def brax_ppo_config(
+    env_name: str, impl: Optional[str] = None
+) -> config_dict.ConfigDict:
   """Returns tuned Brax PPO config for the given environment."""
   env_config = manipulation.get_default_config(env_name)
 
@@ -34,10 +36,11 @@ def brax_ppo_config(env_name: str) -> config_dict.ConfigDict:
           policy_obs_key="state",
           value_obs_key="state",
       ),
+      num_resets_per_eval=10,
   )
   if env_name == "AlohaHandOver":
     rl_config.num_timesteps = 100_000_000
-    rl_config.num_evals = int(rl_config.num_timesteps / 4_000_000)
+    rl_config.num_evals = 25
     rl_config.unroll_length = 15
     rl_config.num_minibatches = 32
     rl_config.num_updates_per_batch = 8
@@ -61,6 +64,9 @@ def brax_ppo_config(env_name: str) -> config_dict.ConfigDict:
     rl_config.num_envs = 1024
     rl_config.batch_size = 512
     rl_config.network_factory.policy_hidden_layer_sizes = (256, 256, 256, 256)
+    if impl == "warp":
+      rl_config.num_timesteps *= 3
+      rl_config.num_evals *= 3
   elif env_name == "PandaOpenCabinet":
     rl_config.num_timesteps = 40_000_000
     rl_config.num_evals = 4
@@ -73,7 +79,6 @@ def brax_ppo_config(env_name: str) -> config_dict.ConfigDict:
     rl_config.num_envs = 2048
     rl_config.batch_size = 512
     rl_config.network_factory.policy_hidden_layer_sizes = (32, 32, 32, 32)
-    rl_config.num_resets_per_eval = 1
   elif env_name == "PandaPickCubeCartesian":
     rl_config.num_timesteps = 5_000_000
     rl_config.num_evals = 5
@@ -89,6 +94,9 @@ def brax_ppo_config(env_name: str) -> config_dict.ConfigDict:
     rl_config.network_factory.policy_hidden_layer_sizes = (256, 256)
     rl_config.num_resets_per_eval = 1
     rl_config.max_grad_norm = 1.0
+    if impl == "warp":
+      rl_config.num_timesteps *= 4
+      rl_config.num_evals *= 4
   elif env_name.startswith("PandaPickCube"):
     rl_config.num_timesteps = 20_000_000
     rl_config.num_evals = 4
@@ -101,6 +109,9 @@ def brax_ppo_config(env_name: str) -> config_dict.ConfigDict:
     rl_config.num_envs = 2048
     rl_config.batch_size = 512
     rl_config.network_factory.policy_hidden_layer_sizes = (32, 32, 32, 32)
+    if impl == "warp":
+      rl_config.num_timesteps *= 4
+      rl_config.num_evals *= 4
   elif env_name == "PandaRobotiqPushCube":
     rl_config.num_timesteps = 1_800_000_000
     rl_config.num_evals = 10
@@ -115,6 +126,10 @@ def brax_ppo_config(env_name: str) -> config_dict.ConfigDict:
     rl_config.num_resets_per_eval = 1
     rl_config.num_eval_envs = 32
     rl_config.network_factory.policy_hidden_layer_sizes = (64, 64, 64, 64)
+    if impl == "warp":
+      rl_config.num_resets_per_eval = 10
+      rl_config.num_timesteps = int(rl_config.num_timesteps * 1.5)
+      rl_config.num_evals = int(rl_config.num_evals * 1.5)
   elif env_name == "LeapCubeRotateZAxis":
     rl_config.num_timesteps = 100_000_000
     rl_config.num_evals = 10
@@ -157,7 +172,9 @@ def brax_ppo_config(env_name: str) -> config_dict.ConfigDict:
   return rl_config
 
 
-def brax_vision_ppo_config(env_name: str) -> config_dict.ConfigDict:
+def brax_vision_ppo_config(
+    env_name: str, unused_impl: Optional[str] = None
+) -> config_dict.ConfigDict:
   """Returns tuned Brax Vision PPO config for the given environment."""
   env_config = manipulation.get_default_config(env_name)
 
@@ -171,6 +188,7 @@ def brax_vision_ppo_config(env_name: str) -> config_dict.ConfigDict:
       network_factory=config_dict.create(
           policy_hidden_layer_sizes=(32, 32, 32, 32)
       ),
+      num_resets_per_eval=10,
   )
 
   if env_name == "PandaPickCubeCartesian":
@@ -192,7 +210,7 @@ def brax_vision_ppo_config(env_name: str) -> config_dict.ConfigDict:
   return rl_config
 
 
-def rsl_rl_config(env_name: str) -> config_dict.ConfigDict:  # pylint: disable=unused-argument
+def rsl_rl_config(env_name: str, unused_impl: Optional[str] = None) -> config_dict.ConfigDict:  # pylint: disable=unused-argument
   """Returns tuned RSL-RL PPO config for the given environment."""
 
   rl_config = config_dict.create(