Update to New RSL-RL - Enable multi gpu

.vscode/settings.json

@@ -0,0 +1,11 @@
+
+{"python.analysis.extraPaths": [
+    "${workspaceFolder}/../../IsaacLab/source/isaaclab",
+    "${workspaceFolder}/../../IsaacLab/source/isaaclab_mimic",
+    "${workspaceFolder}/../../IsaacLab/source/isaaclab_tasks",
+    "${workspaceFolder}/../../IsaacLab/source/isaaclab_assets",
+    "${workspaceFolder}/../../IsaacLab/source/extensions",
+    "${workspaceFolder}/../../IsaacLab/source/standalone",
+    "${workspaceFolder}/../../IsaacLab/source/isaaclab_rl"
+]
+}

README.md

@@ -50,10 +50,6 @@ demonstrations and to link to the original implementation in Isaac Gym, please v
 
 1. Install Isaac Lab, see the [installation guide](https://isaac-sim.github.io/IsaacLab/v2.0.0/source/setup/installation/index.html).
     **Note**: Currently HOVER has been tested with Isaac Lab versions 2.0.0. After you clone the Isaac Lab
-    repository, check out the specific tag before installation. Also note that the `rsl_rl`
-    package is renamed to `rsl_rl_lib` with the current `v2.0.0` tag of Isaac Lab, causing installation issues.
-    This will be fixed once a new tag is created on the Isaac Lab repo.
-    This error would not affect this repo, as we have our own customized `rsl_rl` package.
     ```bash
     git fetch origin
     git checkout v2.0.0
@@ -70,7 +66,7 @@ demonstrations and to link to the original implementation in Isaac Gym, please v
 4. Install this repo and its dependencies by running the following command from the root of this
    repo:
     ```bash
-    ./install_deps.sh
+    pip install -r requirements.txt
     ```
 
 # Training
@@ -121,35 +117,40 @@ folder as the included data library will handle relative path searching, which i
 
 For more details, refer to the [human2humanoid repository](https://github.com/LeCAR-Lab/human2humanoid/tree/main?tab=readme-ov-file#motion-retargeting).
 
+If you have a configured conda environment with Isaaclab you can launch the scripts with `python`, otherwise use the python bundled with Isaaclab `${ISAACLAB_PATH:?}/isaaclab.sh -p`
 ## Teacher Policy
 
 
 In the project's root directory,
 
 ```bash
-${ISAACLAB_PATH:?}/isaaclab.sh -p scripts/rsl_rl/train_teacher_policy.py \
+python scripts/rsl_rl/train_teacher_policy.py \
     --num_envs 1024 \
-    --reference_motion_path neural_wbc/data/data/motions/stable_punch.pkl
+    --reference_motion_path neural_wbc/data/data/motions/stable_punch.pkl # you can omit this as it's the default
 ```
 
-The max iteration of the teacher policy is set to 10,000,000 by default. The resulting checkpoint is stored in `neural_wbc/data/data/policy/h1:teacher/` and the filename is `model_<iteration_number>.pt`.
-The user can interrupt the training at anytime.  Normally, a good policy can be trained between `50k` to `80k` iterations.
+The max iteration of the teacher policy is set to 10,000,000 by default. The resulting checkpoint is stored in `<data_time>` and the filename is `model_<iteration_number>.pt`. The user can interrupt the training at anytime.  Normally, a good policy can be trained between `50k` to `80k` iterations.
 
 
 ## Student Policy
 
 
 In the project's root directory,
 ```bash
-${ISAACLAB_PATH:?}/isaaclab.sh -p scripts/rsl_rl/train_student_policy.py \
+python scripts/rsl_rl/train_student_policy.py \
     --num_envs 1024 \
-    --reference_motion_path neural_wbc/data/data/motions/stable_punch.pkl \
-    --teacher_policy.resume_path neural_wbc/data/data/policy/h1:teacher \
-    --teacher_policy.checkpoint model_<iteration_number>.pt
+    --headless
 ```
 This assumes that you have already trained the teacher policy as there is no provided teacher policy in the repo. Change the filename to match the checkpoint you trained.
 The exact path of the teacher policy does not matter, but it is recommended to store it in the data folder. If stored outside the data folder, you might need to provide the full path.
 
+## Multi-GPU
+It is possible to train on a node with multiple gpus or multiple nodes with multiple gpues.
+In order to train the teacher on a machine with 4 gpus:
+```bash
+python -m torch.distributed.run --nnodes=1 --nproc_per_node=4 scripts/rsl_rl/train_teacher.py --num_envs=1024 --headless --distributed
+```
+This will generate 4 IsaacSim simulation, each containing 1024 environments.
 
 ## General Remarks for Training
 
@@ -159,16 +160,15 @@ The exact path of the teacher policy does not matter, but it is recommended to s
     results we recommend to train with the full amass dataset.
 - Per default the trained checkpoints are stored to `logs/teacher/` or `logs/student/`.
 - If you don't want to train from scratch you can resume training from a checkpoint using the
-    options `--teacher_policy.resume_path`/`--student_policy.resume_path` and
-    `--teacher_policy.checkpoint`/`--student_policy.checkpoint`. For example to resume training of
+    options `--load_run=<date_time>` and/or `--checkpoint=<model_iter.pt>`. For example to resume training of
     the teacher use
 
     ```bash
-    ${ISAACLAB_PATH:?}/isaaclab.sh -p scripts/rsl_rl/train_teacher_policy.py \
+    python scripts/rsl_rl/train_teacher_policy.py \
         --num_envs 10 \
-        --reference_motion_path neural_wbc/data/data/motions/stable_punch.pkl \
-        --teacher_policy.resume_path neural_wbc/data/data/policy/h1:teacher \
-        --teacher_policy.checkpoint model_<iteration_number>.pt
+        --load_run 2025-05-09_09-23-50 \
+        --checkpoint model_500.pt
+
     ```
 
 - Training requires a single GPU, we found the following performance when training on different
@@ -234,24 +234,21 @@ In both cases the same commands from above can be used to launch the training.
 In the project's root directory,
 
 ```bash
-${ISAACLAB_PATH:?}/isaaclab.sh -p scripts/rsl_rl/play.py \
+python scripts/rsl_rl/play.py \
     --num_envs 10 \
     --reference_motion_path neural_wbc/data/data/motions/stable_punch.pkl \
-    --teacher_policy.resume_path neural_wbc/data/data/policy/h1:teacher \
-    --teacher_policy.checkpoint model_<iteration_number>.pt
 ```
+This will run the latest checkpoint in the latest run.
+Alternatively you can specify the run from which you want to load the checkpoint from by passing `--load_run=logs/rsl_rl/teacher_policy/h1/<date_time>`, which will use the latest checkpoint in the folder. You can also specify a specific checkpoint by passing `--checkpoint=<model_iter.pt>`
 
 ## Play Student Policy
 
 In the project's root directory,
 
 ```bash
-${ISAACLAB_PATH:?}/isaaclab.sh -p scripts/rsl_rl/play.py \
+python scripts/rsl_rl/play.py \
     --num_envs 10 \
-    --reference_motion_path neural_wbc/data/data/motions/stable_punch.pkl \
     --student_player \
-    --student_path neural_wbc/data/data/policy/h1:student \
-    --student_checkpoint model_<iteration_number>.pt
 ```
 
 # Evaluation
@@ -297,13 +294,17 @@ The evaluation script, `scripts/rsl_rl/eval.py`, uses the same arguments as the
 `scripts/rsl_rl/play.py`. You can use it for both teacher and student policies.
 
 ```bash
-${ISAACLAB_PATH}/isaaclab.sh -p scripts/rsl_rl/eval.py \
+# For a teacher
+python scripts/rsl_rl/eval.py \
     --num_envs 10 \
-    --teacher_policy.resume_path neural_wbc/data/data/policy/h1:teacher \
-    --teacher_policy.checkpoint model_<iteration_number>.pt
 ```
 
-
+``` bash
+# For a student
+python scripts/rsl_rl/eval.py \
+    --num_envs 10 \
+    --student_player
+```
 # Overwriting Configuration Values
 
 To customize and overwrite default environment configuration values, you can provide a YAML file

neural_wbc/core/neural_wbc/core/environment_wrapper.py

@@ -42,11 +42,15 @@ def reset(self, env_ids: list | torch.Tensor):
         """Resets environment specified by env_ids."""
         raise NotImplementedError
 
-    def get_observations(self) -> torch.Tensor:
+    def get_observations(self) -> tuple[torch.Tensor, dict]:
         """Gets policy observations for each environment based on the mode."""
         if self._mode.is_distill_mode():
-            return self.get_student_observations()
-        return self.get_teacher_observations()
+            return self.get_student_observations(), {
+                "observations": {"policy": self.get_student_observations(), "teacher": self.get_teacher_observations()}
+            }
+        return self.get_teacher_observations(), {
+            "observations": {"policy": self.get_teacher_observations(), "critic": self.get_privileged_observations()}
+        }
 
     def get_teacher_observations(self) -> torch.Tensor:
         """Gets teacher policy observations for each environment."""
@@ -55,3 +59,7 @@ def get_teacher_observations(self) -> torch.Tensor:
     def get_student_observations(self) -> torch.Tensor:
         """Gets student policy observations for each environment."""
         raise NotImplementedError
+
+    def get_privileged_observations(self) -> torch.Tensor:
+        """Gets privileged observations for each environment."""
+        raise NotImplementedError

neural_wbc/core/tests/test_reference_motion_manager.py

@@ -64,10 +64,10 @@ def test_reference_motion_state(self):
         ref_motion_state: ReferenceMotionState = mgr.get_state_from_motion_lib_cache(
             episode_length_buf=episode_length_buf
         )
-        self.assertAlmostEqual(ref_motion_state.body_rot[0, 0, 0], 1)
-        self.assertAlmostEqual(ref_motion_state.body_rot[0, 0, 1], 0)
-        self.assertAlmostEqual(ref_motion_state.body_rot[0, 0, 2], 0)
-        self.assertAlmostEqual(ref_motion_state.body_rot[0, 0, 3], 0)
+        self.assertAlmostEqual(ref_motion_state.body_rot[0, 0, 0].item(), 1)
+        self.assertAlmostEqual(ref_motion_state.body_rot[0, 0, 1].item(), 0)
+        self.assertAlmostEqual(ref_motion_state.body_rot[0, 0, 2].item(), 0)
+        self.assertAlmostEqual(ref_motion_state.body_rot[0, 0, 3].item(), 0)
 
     def test_reset_motion_start_times_after_load_motion(self):
         mgr = self._create_reference_motion_manager()

neural_wbc/inference_env/inference_env/neural_wbc_env.py

@@ -271,7 +271,7 @@ def step(self, actions: torch.Tensor) -> tuple[dict, torch.Tensor, torch.Tensor,
         if self.cfg.mode.is_distill_mode():
             obs = obs_dict["student_policy"]
         else:
-            obs = obs_dict["teacher_policy"]
+            obs = obs_dict["teacher"]
 
         # Extras are required for evaluation
         self.extras["observations"] = obs_dict
@@ -324,7 +324,7 @@ def reset(self, env_ids: list | torch.Tensor | None = None):
         if self.cfg.mode.is_distill_mode():
             self.history.reset(env_ids=env_ids)
 
-        obs = self.get_observations()
+        obs, _ = self.get_observations()
         return obs, None
 
     def get_student_observations(self) -> torch.Tensor:
@@ -336,6 +336,9 @@ def get_student_observations(self) -> torch.Tensor:
         current_obs = self._compute_observations()
         return current_obs["student_policy"]
 
+    def get_teacher_observations(self):
+        pass
+
     def _reset_robot_state_and_motion(self, env_ids: list | torch.Tensor | None = None):
         """Resets the robot state and the reference motion.
 

neural_wbc/isaac_lab_wrapper/neural_wbc/isaac_lab_wrapper/neural_wbc_env.py

@@ -398,7 +398,7 @@ def _get_observations(self) -> dict:
         )
 
         if self.cfg.add_policy_obs_noise:
-            obs_dic["teacher_policy"] = self._observation_noise_model.apply(obs_dic["teacher_policy"])
+            obs_dic["teacher"] = self._observation_noise_model.apply(obs_dic["teacher"])
 
         return obs_dic
 

neural_wbc/isaac_lab_wrapper/neural_wbc/isaac_lab_wrapper/observations.py

@@ -49,7 +49,7 @@ def compute_observations(
         last_actions=env.actions,
     )
     obs_dict.update(teacher_obs_dict)
-    obs_dict["teacher_policy"] = teacher_obs
+    obs_dict["teacher"] = teacher_obs
 
     # Then the privileged observations.
     privileged_obs, privileged_obs_dict = compute_privileged_observations(env=env, asset=asset)
@@ -81,7 +81,7 @@ def compute_privileged_observations(env: NeuralWBCEnv, asset: Articulation):
 
     privileged_obs_dict = {
         "base_com_bias": env.base_com_bias.to(env.device),
-        "ground_friction_values": asset.data.joint_friction[:, env.feet_ids],
+        "ground_friction_values": asset.data.joint_friction_coeff[:, env.feet_ids],
         "body_mass_scale": env.body_mass_scale,
         "kp_scale": env.kp_scale,
         "kd_scale": env.kd_scale,

neural_wbc/isaac_lab_wrapper/tests/test_neural_wbc_env.py

@@ -31,7 +31,7 @@ def _check_observations(self, env: NeuralWBCEnv, obs: dict[str, torch.Tensor]):
             # Make sure that all observations do not contain NaN values.
             self.assertFalse(torch.any(torch.isnan(value)), msg=f"Observation {key} has NaN values {value}.")
         # Ensure that the policy and critic observations have shapes matching the configuration.
-        self.assertEqual(obs["teacher_policy"].shape, env.observation_space.shape)
+        self.assertEqual(obs["teacher"].shape, env.observation_space.shape)
         self.assertEqual(obs["critic"].shape, env.state_space.shape)
 
     def test_step(self):

neural_wbc/isaac_lab_wrapper/tests/test_observations.py

@@ -47,7 +47,7 @@ def setUp(self):
 
         # Mocking the Articulation and its attributes
         self.asset = Mock(spec=Articulation)
-        self.asset.data.joint_friction = torch.randn((self.num_envs, 2))
+        self.asset.data.joint_friction_coeff = torch.randn((self.num_envs, 2))
 
     def test_compute_privileged_observations(self):
         privileged_obs, privileged_obs_dict = compute_privileged_observations(env=self.env, asset=self.asset)
@@ -66,7 +66,8 @@ def test_compute_privileged_observations(self):
         self.assertTrue(torch.equal(privileged_obs_dict["base_com_bias"], self.env.base_com_bias))
         self.assertTrue(
             torch.equal(
-                privileged_obs_dict["ground_friction_values"], self.asset.data.joint_friction[:, self.env.feet_ids]
+                privileged_obs_dict["ground_friction_values"],
+                self.asset.data.joint_friction_coeff[:, self.env.feet_ids],
             )
         )
         self.assertTrue(torch.equal(privileged_obs_dict["body_mass_scale"], self.env.body_mass_scale))

neural_wbc/student_policy/neural_wbc/student_policy/student_policy_trainer.py

@@ -99,9 +99,9 @@ def run(self):
                 with torch.inference_mode():
                     actions = self._produce_actions(obs_dict)
                     obs, privileged_obs, rewards, dones, infos = self._env.step(actions.detach())
-                    gt_actions = self._get_ground_truth_actions(obs_dict["teacher_policy"])
+                    gt_actions = self._get_ground_truth_actions(obs_dict["teacher"])
                     slice = Slice(
-                        policy_observations=obs_dict["teacher_policy"],
+                        policy_observations=obs_dict["teacher"],
                         student_observations=obs_dict["student_policy"],
                         ground_truth_actions=gt_actions,
                         applied_actions=actions,
@@ -189,7 +189,7 @@ def log(self, locs, width=80, pad=35):
 
         str = (
             " \033[1m Learning iteration"
-            f" {self._iterations+self.start_iteration}/{self._cfg.max_iteration+self.start_iteration} \033[0m "
+            f" {self._iterations + self.start_iteration} / {self._cfg.max_iteration + self.start_iteration} \033[0m "
         )
 
         if len(locs["rewbuffer"]) > 0:
@@ -252,7 +252,7 @@ def _produce_actions(self, obs: dict) -> torch.Tensor:
             observations = obs["student_policy"]
             action = self._student.act(observations)
         else:
-            action = self._teacher.act_rollout(obs["teacher_policy"])
+            action = self._teacher.act_rollout(obs["teacher"])
 
         return action
 

neural_wbc/student_policy/neural_wbc/tests/test_student_policy_trainer.py

@@ -59,11 +59,11 @@ def __init__(self):
 
     def _generate_random_observations(self):
         obs = {
-            "teacher_policy": torch.randn(self.num_envs, self.num_obs, device=self.device),
+            "teacher": torch.randn(self.num_envs, self.num_obs, device=self.device),
             "critic": torch.randn(self.num_envs, self.num_privileged_obs, device=self.device),
             "student_policy": torch.randn(self.num_envs, self.num_student_obs, device=self.device),
         }
-        self.obs_buf = obs["teacher_policy"]
+        self.obs_buf = obs["teacher"]
         self.privileged_obs_buf = obs["critic"]
         self.student_obs_buf = obs["student_policy"]
         return obs
@@ -139,7 +139,7 @@ def test_run(self):
         config_path = os.path.join(self.student_policy_path, "config.json")
         with open(config_path) as fh:
             config_dict = json.load(fh)
-        config_dict["teacher_policy"] = teacher_policy
+        config_dict["teacher"] = teacher_policy
         reconstructed_cfg = StudentPolicyTrainerCfg(**config_dict)
         self.assertEqual(cfg, reconstructed_cfg)
 

pyproject.toml

@@ -6,6 +6,9 @@ build-backend = "setuptools.build_meta"
 name = "neural_wbc"
 version = "0.1.0"
 
+[tool.setuptools.packages.find]
+include = ["neural_wbc"]
+
 [tool.isort]
 atomic = true
 profile = "black"
@@ -59,10 +62,6 @@ known_isaaclabparty = [
     "isaaclab_assets"
 ]
 
-[tool.setuptools]
-package-dir = {"rsl_rl" = "third_party/rsl_rl/rsl_rl"}
-packages = ["rsl_rl"]
-
 [tool.pyright]
 exclude = [
     "**/__pycache__",

requirements.txt

@@ -4,7 +4,7 @@ git+https://github.com/ZhengyiLuo/SMPLSim.git@dd65a86
 easydict
 warp-lang
 dataclass-wizard
-
+git+https://github.com/leggedrobotics/rsl_rl.git@main
 -e neural_wbc/core
 -e neural_wbc/data
 -e neural_wbc/isaac_lab_wrapper
@@ -13,4 +13,3 @@ dataclass-wizard
 -e neural_wbc/student_policy
 -e third_party/human2humanoid/phc
 -e third_party/mujoco_viewer
--e third_party/rsl_rl