Update manipulation environments for Warp backend.

PiperOrigin-RevId: 831760691
Change-Id: I63716e9549a624bb6b4782687ddff515f17243eb

Author: thowell

mujoco_playground/_src/manipulation/aloha/handover.py

@@ -42,7 +42,7 @@ def default_config() -> config_dict.ConfigDict:
           ),
       ),
       impl='jax',
-      nconmax=24 * 8192,
+      nconmax=24 * 2048,
       njmax=88,
   )
 

mujoco_playground/_src/manipulation/aloha/single_peg_insertion.py

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

mujoco_playground/_src/manipulation/franka_emika_panda/open_cabinet.py

@@ -49,7 +49,7 @@ def default_config() -> config_dict.ConfigDict:
           )
       ),
       impl="jax",
-      nconmax=12 * 8192,
+      nconmax=12 * 2048,
       njmax=96,
   )
 

mujoco_playground/_src/manipulation/franka_emika_panda/pick.py

@@ -48,7 +48,7 @@ def default_config() -> config_dict.ConfigDict:
           )
       ),
       impl='jax',
-      nconmax=24 * 8192,
+      nconmax=24 * 2048,
       njmax=128,
   )
   return config

mujoco_playground/_src/manipulation/franka_emika_panda/pick_cartesian.py

@@ -75,7 +75,7 @@ def default_config():
       success_threshold=0.05,
       action_history_length=1,
       impl='jax',
-      nconmax=12 * 8192,
+      nconmax=12 * 1024,
       njmax=128,
   )
   return config

mujoco_playground/config/manipulation_params.py

@@ -64,9 +64,6 @@ def brax_ppo_config(
     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
@@ -94,9 +91,6 @@ def brax_ppo_config(
     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
@@ -109,9 +103,6 @@ def brax_ppo_config(
     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
@@ -126,10 +117,6 @@ def brax_ppo_config(
     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