internal change

PiperOrigin-RevId: 832442898
Change-Id: Idf68bc40b5d63c29681a5852599967e2d149a22c

Author: btaba

learning/train_rsl_rl.py

@@ -15,15 +15,9 @@
 # pylint: disable=wrong-import-position
 """Train a PPO agent using RSL-RL for the specified environment."""
 
-import os
-
-xla_flags = os.environ.get("XLA_FLAGS", "")
-xla_flags += " --xla_gpu_triton_gemm_any=True"
-os.environ["XLA_FLAGS"] = xla_flags
-os.environ["MUJOCO_GL"] = "egl"
-
 from datetime import datetime
 import json
+import os
 
 from absl import app
 from absl import flags
@@ -32,15 +26,24 @@
 import mediapy as media
 from ml_collections import config_dict
 import mujoco
-from rsl_rl.runners import OnPolicyRunner
-import torch
-import wandb
-
 import mujoco_playground
 from mujoco_playground import registry
 from mujoco_playground import wrapper_torch
 from mujoco_playground.config import locomotion_params
 from mujoco_playground.config import manipulation_params
+from rsl_rl.runners import OnPolicyRunner
+import torch
+import warp as wp
+
+try:
+  import wandb  # pylint: disable=g-import-not-at-top
+except ImportError:
+  wandb = None
+
+xla_flags = os.environ.get("XLA_FLAGS", "")
+xla_flags += " --xla_gpu_triton_gemm_any=True"
+os.environ["XLA_FLAGS"] = xla_flags
+os.environ["MUJOCO_GL"] = "egl"
 
 # Suppress logs if you want
 logging.set_verbosity(logging.WARNING)
@@ -78,6 +81,11 @@
 _CAMERA = flags.DEFINE_string(
     "camera", None, "Camera name to use for rendering."
 )
+_WP_KERNEL_CACHE_DIR = flags.DEFINE_string(
+    "wp_kernel_cache_dir",
+    "/tmp/wp_kernel_cache_playground",
+    "Path to the WP kernel cache directory.",
+)
 
 
 def get_rl_config(env_name: str) -> config_dict.ConfigDict:
@@ -93,6 +101,8 @@ def main(argv):
   """Run training and evaluation for the specified environment using RSL-RL."""
   del argv  # unused
 
+  wp.config.kernel_cache_dir = _WP_KERNEL_CACHE_DIR.value
+
   # Possibly parse the device for multi-GPU
   if _MULTI_GPU.value:
     local_rank = int(os.environ.get("LOCAL_RANK", "0"))
@@ -119,7 +129,7 @@ def main(argv):
   print(f"Experiment name: {exp_name}")
 
   # Logging directory
-  logdir = os.path.abspath(os.path.join("logs", exp_name))
+  logdir = os.path.abspath(os.path.join("/tmp/rslrl-training-logs/", exp_name))
   os.makedirs(logdir, exist_ok=True)
   print(f"Logs are being stored in: {logdir}")
 
@@ -129,7 +139,7 @@ def main(argv):
   print(f"Checkpoint path: {ckpt_path}")
 
   # Initialize Weights & Biases if required
-  if _USE_WANDB.value and not _PLAY_ONLY.value:
+  if _USE_WANDB.value and not _PLAY_ONLY.value and wandb is not None:
     wandb.tensorboard.patch(root_logdir=logdir)
     wandb.init(project="mjxrl", name=exp_name)
     wandb.config.update(env_cfg.to_dict())
@@ -152,7 +162,9 @@ def render_callback(_, state):
     render_trajectory.append(state)
 
   # Create the environment
-  raw_env = registry.load(_ENV_NAME.value, config=env_cfg)
+  raw_env = registry.load(
+      _ENV_NAME.value, config=env_cfg, config_overrides={"impl": "jax"}
+  )
   brax_env = wrapper_torch.RSLRLBraxWrapper(
       raw_env,
       num_envs,
@@ -186,7 +198,7 @@ def render_callback(_, state):
   # If resume, load from checkpoint
   if train_cfg.resume:
     resume_path = wrapper_torch.get_load_path(
-        os.path.abspath("logs"),
+        "/tmp/rslrl-training-logs/",
         load_run=train_cfg.load_run,
         checkpoint=train_cfg.checkpoint,
     )
@@ -206,7 +218,9 @@ def render_callback(_, state):
   policy = runner.get_inference_policy(device=device)
 
   # Example: run a single rollout
-  eval_env = registry.load(_ENV_NAME.value, config=env_cfg)
+  eval_env = registry.load(
+      _ENV_NAME.value, config=env_cfg, config_overrides={"impl": "jax"}
+  )
   jit_reset = jax.jit(eval_env.reset)
   jit_step = jax.jit(eval_env.step)
 
@@ -215,18 +229,25 @@ def render_callback(_, state):
   rollout = [state]
 
   # We’ll assume your environment’s observation is in state.obs["state"].
-  obs_torch = wrapper_torch._jax_to_torch(state.obs["state"])
+  is_dict_obs = isinstance(eval_env.observation_size, dict)
+  obs = state.obs["state"] if is_dict_obs else state.obs
+  obs_torch = wrapper_torch._jax_to_torch(obs)
 
   for _ in range(env_cfg.episode_length):
     with torch.no_grad():
-      actions = policy(obs_torch)
+      actions = policy({"state": obs_torch})
+      actions = torch.clip(actions, -1.0, 1.0)  # from wrapper_torch.py
     # Step environment
     state = jit_step(state, wrapper_torch._torch_to_jax(actions.flatten()))
     rollout.append(state)
-    obs_torch = wrapper_torch._jax_to_torch(state.obs["state"])
+    obs = state.obs["state"] if is_dict_obs else state.obs
+    obs_torch = wrapper_torch._jax_to_torch(obs)
     if state.done:
       break
 
+  reward_sum = sum(s.reward for s in rollout)
+  print(f"Rollout reward: {reward_sum}")
+
   # Render
   scene_option = mujoco.MjvOption()
   scene_option.flags[mujoco.mjtVisFlag.mjVIS_TRANSPARENT] = True

mujoco_playground/config/manipulation_params.py

@@ -216,18 +216,18 @@ def rsl_rl_config(env_name: str, unused_impl: Optional[str] = None) -> config_di
           value_loss_coef=1.0,
           use_clipped_value_loss=True,
           clip_param=0.2,
-          entropy_coef=0.001,
-          num_learning_epochs=5,
+          entropy_coef=0.01,
+          num_learning_epochs=4,
           # mini batch size = num_envs*nsteps / nminibatches
-          num_mini_batches=4,
-          learning_rate=3.0e-4,  # 5.e-4
+          num_mini_batches=8,
+          learning_rate=1e-3,
           schedule="adaptive",  # could be adaptive, fixed
-          gamma=0.99,
+          gamma=0.97,
           lam=0.95,
           desired_kl=0.01,
           max_grad_norm=1.0,
       ),
-      num_steps_per_env=24,  # per iteration
+      num_steps_per_env=40,  # per iteration
       max_iterations=100000,  # number of policy updates
       empirical_normalization=True,
       # logging