diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 0000000..522b04f --- /dev/null +++ b/.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" +] +} diff --git a/README.md b/README.md index e993759..6971d1a 100644 --- a/README.md +++ b/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,19 +117,19 @@ 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_.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 `` and the filename is `model_.pt`. The user can interrupt the training at anytime. Normally, a good policy can be trained between `50k` to `80k` iterations. ## Student Policy @@ -141,15 +137,20 @@ The user can interrupt the training at anytime. Normally, a good policy can be 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_.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=` and/or `--checkpoint=`. 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_.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_.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/`, which will use the latest checkpoint in the folder. You can also specify a specific checkpoint by passing `--checkpoint=` ## 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_.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_.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 diff --git a/neural_wbc/core/neural_wbc/core/environment_wrapper.py b/neural_wbc/core/neural_wbc/core/environment_wrapper.py index 1922b53..e74ecdc 100644 --- a/neural_wbc/core/neural_wbc/core/environment_wrapper.py +++ b/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 diff --git a/neural_wbc/core/tests/test_reference_motion_manager.py b/neural_wbc/core/tests/test_reference_motion_manager.py index d6d4aae..aaf87ec 100644 --- a/neural_wbc/core/tests/test_reference_motion_manager.py +++ b/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() diff --git a/neural_wbc/data/data/motions/stable_punch.pkl b/neural_wbc/data/data/motions/stable_punch.pkl new file mode 100644 index 0000000..d98a7b8 Binary files /dev/null and b/neural_wbc/data/data/motions/stable_punch.pkl differ diff --git a/neural_wbc/inference_env/inference_env/neural_wbc_env.py b/neural_wbc/inference_env/inference_env/neural_wbc_env.py index 96f6364..77dd28d 100644 --- a/neural_wbc/inference_env/inference_env/neural_wbc_env.py +++ b/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. diff --git a/neural_wbc/isaac_lab_wrapper/neural_wbc/isaac_lab_wrapper/neural_wbc_env.py b/neural_wbc/isaac_lab_wrapper/neural_wbc/isaac_lab_wrapper/neural_wbc_env.py index a6e5006..a6f5f1c 100644 --- a/neural_wbc/isaac_lab_wrapper/neural_wbc/isaac_lab_wrapper/neural_wbc_env.py +++ b/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 diff --git a/neural_wbc/isaac_lab_wrapper/neural_wbc/isaac_lab_wrapper/observations.py b/neural_wbc/isaac_lab_wrapper/neural_wbc/isaac_lab_wrapper/observations.py index cf497b1..de6a1ac 100644 --- a/neural_wbc/isaac_lab_wrapper/neural_wbc/isaac_lab_wrapper/observations.py +++ b/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, diff --git a/neural_wbc/isaac_lab_wrapper/tests/test_neural_wbc_env.py b/neural_wbc/isaac_lab_wrapper/tests/test_neural_wbc_env.py index ad71d19..e542367 100644 --- a/neural_wbc/isaac_lab_wrapper/tests/test_neural_wbc_env.py +++ b/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): diff --git a/neural_wbc/isaac_lab_wrapper/tests/test_observations.py b/neural_wbc/isaac_lab_wrapper/tests/test_observations.py index 1a7aeb5..e15b2ac 100644 --- a/neural_wbc/isaac_lab_wrapper/tests/test_observations.py +++ b/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)) diff --git a/neural_wbc/student_policy/neural_wbc/student_policy/student_policy_trainer.py b/neural_wbc/student_policy/neural_wbc/student_policy/student_policy_trainer.py index def2791..83f0b4c 100644 --- a/neural_wbc/student_policy/neural_wbc/student_policy/student_policy_trainer.py +++ b/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 diff --git a/neural_wbc/student_policy/neural_wbc/tests/test_student_policy_trainer.py b/neural_wbc/student_policy/neural_wbc/tests/test_student_policy_trainer.py index e749334..0eb24b8 100644 --- a/neural_wbc/student_policy/neural_wbc/tests/test_student_policy_trainer.py +++ b/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) diff --git a/pyproject.toml b/pyproject.toml index 6f539c7..083ddf9 100644 --- a/pyproject.toml +++ b/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__", diff --git a/requirements.txt b/requirements.txt index 1f1849d..b5ed813 100644 --- a/requirements.txt +++ b/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 diff --git a/scripts/rsl_rl/cli_args.py b/scripts/rsl_rl/cli_args.py new file mode 100644 index 0000000..96747ae --- /dev/null +++ b/scripts/rsl_rl/cli_args.py @@ -0,0 +1,91 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import argparse +import random +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg + + +def add_rsl_rl_args(parser: argparse.ArgumentParser): + """Add RSL-RL arguments to the parser. + + Args: + parser: The parser to add the arguments to. + """ + # create a new argument group + arg_group = parser.add_argument_group("rsl_rl", description="Arguments for RSL-RL agent.") + # -- experiment arguments + arg_group.add_argument( + "--experiment_name", type=str, default=None, help="Name of the experiment folder where logs will be stored." + ) + arg_group.add_argument("--run_name", type=str, default=None, help="Run name suffix to the log directory.") + # -- load arguments + arg_group.add_argument("--resume", type=bool, default=None, help="Whether to resume from a checkpoint.") + arg_group.add_argument("--load_run", type=str, default=None, help="Name of the run folder to resume from.") + arg_group.add_argument("--checkpoint", type=str, default=None, help="Checkpoint file to resume from.") + # -- logger arguments + arg_group.add_argument( + "--logger", type=str, default=None, choices={"wandb", "tensorboard", "neptune"}, help="Logger module to use." + ) + arg_group.add_argument( + "--log_project_name", type=str, default=None, help="Name of the logging project when using wandb or neptune." + ) + + +def parse_rsl_rl_cfg(task_name: str, args_cli: argparse.Namespace) -> RslRlOnPolicyRunnerCfg: + """Parse configuration for RSL-RL agent based on inputs. + + Args: + task_name: The name of the environment. + args_cli: The command line arguments. + + Returns: + The parsed configuration for RSL-RL agent based on inputs. + """ + from isaaclab_tasks.utils.parse_cfg import load_cfg_from_registry + + # load the default configuration + rslrl_cfg: RslRlOnPolicyRunnerCfg = load_cfg_from_registry(task_name, "rsl_rl_cfg_entry_point") + rslrl_cfg = update_rsl_rl_cfg(rslrl_cfg, args_cli) + return rslrl_cfg + + +def update_rsl_rl_cfg(agent_cfg: RslRlOnPolicyRunnerCfg, args_cli: argparse.Namespace): + """Update configuration for RSL-RL agent based on inputs. + + Args: + agent_cfg: The configuration for RSL-RL agent. + args_cli: The command line arguments. + + Returns: + The updated configuration for RSL-RL agent based on inputs. + """ + # override the default configuration with CLI arguments + if hasattr(args_cli, "seed") and args_cli.seed is not None: + # randomly sample a seed if seed = -1 + if args_cli.seed == -1: + args_cli.seed = random.randint(0, 10000) + agent_cfg.seed = args_cli.seed + if args_cli.resume is not None: + agent_cfg.resume = args_cli.resume + if args_cli.load_run is not None: + agent_cfg.load_run = args_cli.load_run + if args_cli.checkpoint is not None: + agent_cfg.load_checkpoint = args_cli.checkpoint + if args_cli.run_name is not None: + agent_cfg.run_name = args_cli.run_name + if args_cli.logger is not None: + agent_cfg.logger = args_cli.logger + # set the project name for wandb and neptune + if agent_cfg.logger in {"wandb", "neptune"} and args_cli.log_project_name: + agent_cfg.wandb_project = args_cli.log_project_name + agent_cfg.neptune_project = args_cli.log_project_name + + return agent_cfg diff --git a/scripts/rsl_rl/eval.py b/scripts/rsl_rl/eval.py index 7fed252..e0b34f0 100644 --- a/scripts/rsl_rl/eval.py +++ b/scripts/rsl_rl/eval.py @@ -17,7 +17,7 @@ import pprint import yaml -from teacher_policy_cfg import TeacherPolicyCfg +import cli_args from isaaclab.app import AppLauncher @@ -30,9 +30,9 @@ parser.add_argument( "--env_config_overwrite", type=str, default=None, help="Path to yaml file with overwriting configuration values." ) +cli_args.add_rsl_rl_args(parser) + -# append RSL-RL cli arguments -TeacherPolicyCfg.add_args_to_parser(parser=parser) # append AppLauncher cli args AppLauncher.add_app_launcher_args(parser) args_cli = parser.parse_args() diff --git a/scripts/rsl_rl/play.py b/scripts/rsl_rl/play.py index f378be2..3c1b737 100644 --- a/scripts/rsl_rl/play.py +++ b/scripts/rsl_rl/play.py @@ -13,21 +13,18 @@ # See the License for the specific language governing permissions and # limitations under the License. - -from teacher_policy_cfg import TeacherPolicyCfg - from isaaclab.app import AppLauncher # local imports from utils import get_player_args # isort: skip +import cli_args # isort: skip # add argparse arguments parser = get_player_args(description="Plays motion tracking policy in Isaac Lab.") parser.add_argument("--randomize", action="store_true", help="Whether to randomize reference motion while playing.") +cli_args.add_rsl_rl_args(parser) -# append RSL-RL cli arguments -TeacherPolicyCfg.add_args_to_parser(parser) # append AppLauncher cli args AppLauncher.add_app_launcher_args(parser) args_cli = parser.parse_args() @@ -36,6 +33,7 @@ app_launcher = AppLauncher(args_cli) simulation_app = app_launcher.app + from players import DemoPlayer diff --git a/scripts/rsl_rl/players.py b/scripts/rsl_rl/players.py index b006b30..c7d7fde 100644 --- a/scripts/rsl_rl/players.py +++ b/scripts/rsl_rl/players.py @@ -15,22 +15,24 @@ import argparse -import json import os import pprint import torch from typing import Any -from isaaclab_rl.rsl_rl import export_policy_as_onnx -from teacher_policy_cfg import TeacherPolicyCfg -from utils import get_ppo_runner_and_checkpoint_path +import cli_args +from isaaclab_rl.rsl_rl import export_policy_as_jit, export_policy_as_onnx +from rsl_rl.runners import OnPolicyRunner +from student_policy_cfg import StudentPolicyRunnerCfg +from teacher_policy_cfg import TeacherPolicyRunnerCfg from vecenv_wrapper import RslRlNeuralWBCVecEnvWrapper from neural_wbc.core.evaluator import Evaluator from neural_wbc.core.modes import NeuralWBCModes from neural_wbc.isaac_lab_wrapper.neural_wbc_env import NeuralWBCEnv from neural_wbc.isaac_lab_wrapper.neural_wbc_env_cfg_h1 import NeuralWBCEnvCfgH1 -from neural_wbc.student_policy import StudentPolicyTrainer, StudentPolicyTrainerCfg + +from isaaclab_tasks.utils import get_checkpoint_path class Player: @@ -57,35 +59,36 @@ def __init__(self, args_cli: argparse.Namespace, randomize: bool, custom_config: # Create environment and wrap it for RSL RL. self.env = NeuralWBCEnv(cfg=env_cfg) - self.wrapped_env = RslRlNeuralWBCVecEnvWrapper(self.env) - - if self.student_player: - student_path = args_cli.student_path - if student_path: - with open(os.path.join(student_path, "config.json")) as fh: - config_dict = json.load(fh) - config_dict["teacher_policy"] = None - config_dict["resume_path"] = student_path - config_dict["checkpoint"] = args_cli.student_checkpoint - student_cfg = StudentPolicyTrainerCfg(**config_dict) - student_trainer = StudentPolicyTrainer(env=self.wrapped_env, cfg=student_cfg) - self.policy = student_trainer.get_inference_policy(device=self.env.device) - else: - raise ValueError("student_policy.resume_path is needed for play or eval. Please specify a value.") - else: - teacher_policy_cfg = TeacherPolicyCfg.from_argparse_args(args_cli) - ppo_runner, checkpoint_path = get_ppo_runner_and_checkpoint_path( - teacher_policy_cfg=teacher_policy_cfg, wrapped_env=self.wrapped_env, device=self.env.device - ) - ppo_runner.load(checkpoint_path) - print(f"[INFO]: Loaded model checkpoint from: {checkpoint_path}") - - # obtain the trained policy for inference - self.policy = ppo_runner.get_inference_policy(device=self.env.device) - - # export policy to onnx - export_model_dir = os.path.join(os.path.dirname(checkpoint_path), "exported") - export_policy_as_onnx(ppo_runner.alg.actor_critic, export_model_dir, filename="policy.onnx") + self.env = RslRlNeuralWBCVecEnvWrapper(self.env) + + runner_cfg = TeacherPolicyRunnerCfg() if not self.student_player else StudentPolicyRunnerCfg() + agent_cfg = cli_args.update_rsl_rl_cfg(runner_cfg, args_cli) + + log_root_path = os.path.join("logs", "rsl_rl", args_cli.robot, agent_cfg.experiment_name) + log_root_path = os.path.abspath(log_root_path) + + # load the policy from checkpoint + print(f"[INFO] Loading experiment from directory: {log_root_path}") + resume_path = get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint) + + ppo_runner = OnPolicyRunner(env=self.env, train_cfg=agent_cfg.to_dict(), log_dir=None) + ppo_runner.load(resume_path) + print(f"[INFO]: Loaded model checkpoint from: {resume_path}") + + # obtain the trained policy for inference + self.policy = ppo_runner.get_inference_policy(device=self.env.device) + + # export policy to onnx + export_model_dir = os.path.join(os.path.dirname(resume_path), "exported") + export_policy_as_jit( + ppo_runner.alg.policy, ppo_runner.obs_normalizer, path=export_model_dir, filename="policy.pt" + ) + export_policy_as_onnx( + ppo_runner.alg.policy, + normalizer=ppo_runner.obs_normalizer, + path=export_model_dir, + filename="policy.onnx", + ) def _update_env_cfg(self, env_cfg: NeuralWBCEnvCfgH1, custom_config: dict[str, Any]): for key, value in custom_config.items(): @@ -101,7 +104,7 @@ def _update_env_cfg(self, env_cfg: NeuralWBCEnvCfgH1, custom_config: dict[str, A pprint.pprint(env_cfg) def play(self, simulation_app): - obs = self.wrapped_env.get_observations() + obs, extras = self.env.get_observations() # simulate environment while simulation_app.is_running() and not self._should_stop(): @@ -110,9 +113,13 @@ def play(self, simulation_app): # agent stepping actions = self.policy(obs) # env stepping - obs, privileged_obs, rewards, dones, extras = self.wrapped_env.step(actions) + obs, rewards, dones, extras = self.env.step(actions) obs = self._post_step( - obs=obs, privileged_obs=privileged_obs, rewards=rewards, dones=dones, extras=extras + obs=obs, + privileged_obs=extras["observations"]["critic"], + rewards=rewards, + dones=dones, + extras=extras, ) # close the simulator @@ -149,7 +156,7 @@ def __init__( self, args_cli: argparse.Namespace, metrics_path: str | None = None, custom_config: dict[str, Any] | None = None ): super().__init__(randomize=False, args_cli=args_cli, custom_config=custom_config) - self._evaluator = Evaluator(env_wrapper=self.wrapped_env, metrics_path=metrics_path) + self._evaluator = Evaluator(env_wrapper=self.env, metrics_path=metrics_path) def play(self, simulation_app): super().play(simulation_app=simulation_app) @@ -164,6 +171,6 @@ def _post_step( reset_env = self._evaluator.collect(dones=dones, info=extras) if reset_env and not self._evaluator.is_evaluation_complete(): self._evaluator.forward_motion_samples() - obs, _ = self.wrapped_env.reset() + obs, _ = self.env.reset() return obs diff --git a/scripts/rsl_rl/student_policy_cfg.py b/scripts/rsl_rl/student_policy_cfg.py new file mode 100644 index 0000000..cca33f4 --- /dev/null +++ b/scripts/rsl_rl/student_policy_cfg.py @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import MISSING + +from isaaclab_rl.rsl_rl import RslRlDistillationAlgorithmCfg, RslRlDistillationStudentTeacherCfg, RslRlOnPolicyRunnerCfg + +from isaaclab.utils import configclass + + +@configclass +class RslDistillationAlgorithm_Extended_Cfg(RslRlDistillationAlgorithmCfg): + """Configuration settings for the distillation algorithm.""" + + # Add max_grad_norm to the configuration + # As it has not been added to isaaclab yet + max_grad_norm: float = MISSING + pass + + +@configclass +class StudentPolicyRunnerCfg(RslRlOnPolicyRunnerCfg): + """Configuration settings for the runner.""" + + num_steps_per_env = 1 + max_iterations = 200000 + save_interval = 500 + experiment_name = "student_policy" + empirical_normalization = False + policy = RslRlDistillationStudentTeacherCfg( + student_hidden_dims=[512, 256, 128], + teacher_hidden_dims=[512, 256, 128], + activation="elu", + init_noise_std=0.001, + ) + algorithm = RslDistillationAlgorithm_Extended_Cfg( + num_learning_epochs=5, + learning_rate=5e-04, + gradient_length=1, + max_grad_norm=0.2, + ) diff --git a/scripts/rsl_rl/teacher_policy_cfg.py b/scripts/rsl_rl/teacher_policy_cfg.py index bc7ded5..9eaf831 100644 --- a/scripts/rsl_rl/teacher_policy_cfg.py +++ b/scripts/rsl_rl/teacher_policy_cfg.py @@ -12,180 +12,38 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -import argparse -import json -from dataclasses import MISSING, asdict, dataclass, field, fields, is_dataclass -from typing import List, Type -from dataclass_wizard import fromdict +from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlPpoActorCriticCfg, RslRlPpoAlgorithmCfg +from isaaclab.utils import configclass -@dataclass -class PolicyCfg: - """Configuration settings for the policy.""" - ACTIVATION_TYPES = ["elu", "selu", "relu", "crelu", "lrelu", "tanh", "sigmoid"] - - init_noise_std: float = field(default=1.0, metadata={"description": "Initial noise standard deviation for policy."}) - actor_hidden_dims: List[int] = field( - default_factory=lambda: [512, 256, 128], metadata={"description": "Hidden dimensions for the actor network."} - ) - critic_hidden_dims: List[int] = field( - default_factory=lambda: [512, 256, 128], metadata={"description": "Hidden dimensions for the critic network."} - ) - activation: str = field(default="elu", metadata={"description": "Activation function for the neural networks."}) - - def __post_init__(self): - if self.activation not in self.ACTIVATION_TYPES: - raise ValueError( - f"Invalid policy activation {self.activation}. Allowed values are {self.ACTIVATION_TYPES}." - ) - - -@dataclass -class AlgorithmCfg: - """Configuration settings for the algorithm.""" - - value_loss_coef: float = field(default=1.0, metadata={"description": "Coefficient for value loss."}) - use_clipped_value_loss: bool = field(default=True, metadata={"description": "Whether to use clipped value loss."}) - clip_param: float = field(default=0.2, metadata={"description": "Clipping parameter for PPO."}) - entropy_coef: float = field(default=0.005, metadata={"description": "Coefficient for entropy regularization."}) - num_learning_epochs: int = field(default=5, metadata={"description": "Number of learning epochs per update."}) - num_mini_batches: int = field(default=4, metadata={"description": "Number of mini-batches per epoch."}) - learning_rate: float = field(default=1.0e-3, metadata={"description": "Learning rate for the optimizer."}) - schedule: str = field(default="adaptive", metadata={"description": "Learning rate schedule type."}) - gamma: float = field(default=0.99, metadata={"description": "Discount factor for rewards."}) - lam: float = field(default=0.95, metadata={"description": "Lambda for Generalized Advantage Estimation."}) - desired_kl: float = field(default=0.01, metadata={"description": "Desired KL divergence for policy updates."}) - max_grad_norm: float = field(default=0.2, metadata={"description": "Maximum norm for gradient clipping."}) - - -@dataclass -class RunnerCfg: +@configclass +class TeacherPolicyRunnerCfg(RslRlOnPolicyRunnerCfg): """Configuration settings for the runner.""" - path: str = field(default="logs/teacher", metadata={"description": "Directory to save policy and logs to"}) - policy_class_name: str = field(default="ActorCritic", metadata={"description": "Class name for the policy."}) - algorithm_class_name: str = field(default="PPO", metadata={"description": "Class name for the algorithm."}) - num_steps_per_env: int = field( - default=24, metadata={"description": "Number of steps per environment per iteration."} + num_steps_per_env = 24 + max_iterations = 10000000 + save_interval = 500 + experiment_name = "teacher_policy" + empirical_normalization = False + policy = RslRlPpoActorCriticCfg( + actor_hidden_dims=[512, 256, 128], + critic_hidden_dims=[512, 256, 128], + activation="elu", + init_noise_std=1.0, ) - max_iterations: int = field(default=10000000, metadata={"description": "Maximum number of policy updates."}) - save_interval: int = field(default=500, metadata={"description": "Interval for saving checkpoints."}) - checkpoint: str = field(default="", metadata={"description": "Checkpoint to load."}) - resume: bool = field(default=False, metadata={"description": "Whether to resume from a checkpoint."}) - resume_path: str = field(default="", metadata={"description": "Path to the checkpoint for resuming."}) - - @staticmethod - def args_prefix(): - """ - Get the prefix to be used for command-line arguments. - - Returns: - str: The prefix string for command-line arguments. - """ - return "runner." - - -@dataclass -class TeacherPolicyCfg: - """Main configuration class that aggregates all required configurations for teacher policy training.""" - - seed: int = field(default=1, metadata={"description": "Random seed for reproducibility."}) - policy: PolicyCfg = field(default_factory=PolicyCfg, metadata={"description": "Policy configuration settings."}) - algorithm: AlgorithmCfg = field( - default_factory=AlgorithmCfg, metadata={"description": "Algorithm configuration settings."} + algorithm = RslRlPpoAlgorithmCfg( + value_loss_coef=1.0, + use_clipped_value_loss=True, + clip_param=0.2, + entropy_coef=0.005, + num_learning_epochs=5, + num_mini_batches=4, + learning_rate=1.0e-3, + schedule="adaptive", + gamma=0.99, + lam=0.95, + desired_kl=0.01, + max_grad_norm=1.0, ) - runner: RunnerCfg = field(default_factory=RunnerCfg, metadata={"description": "Runner configuration settings."}) - - def to_dict(self): - return asdict(self) - - def save(self, filepath: str): - """Saves the configuration to a YAML file.""" - with open(filepath, "w", encoding="utf-8") as file: - json.dump(self.to_dict(), file) - - @staticmethod - def load(filepath: str) -> "TeacherPolicyCfg": - """Loads the configuration from a YAML file.""" - with open(filepath, encoding="utf-8") as file: - data = json.load(file) - cfg = fromdict(TeacherPolicyCfg, data) - return cfg - - @staticmethod - def args_prefix(): - """ - Get the prefix to be used for command-line arguments. - - Returns: - str: The prefix string for command-line arguments. - """ - return "teacher_policy." - - @staticmethod - def add_args_to_parser(parser: argparse.ArgumentParser, default_overwrites: dict = {}): - """Adds configuration fields to an ArgumentParser.""" - group = parser.add_argument_group("Teacher policy configurations (RSL RL)") - - def add_fields_to_parser(fields): - for field_info in fields: - arg_name = f"--{TeacherPolicyCfg.args_prefix()}{field_info.name}" - arg_type = field_info.type - arg_help = field_info.metadata.get("description", "") - default = None - - if field_info.name in default_overwrites: - default = default_overwrites[field_info.name] - elif field_info.default is not MISSING: - default = field_info.default - elif field_info.default_factory is not MISSING: - default = field_info.default_factory() - - if isinstance(default, list): # Special handling for lists - group.add_argument(arg_name, type=type(default[0]), nargs="+", default=default, help=arg_help) - elif isinstance(default, bool): - if default: - group.add_argument(arg_name, action="store_false", help=arg_help) - else: - group.add_argument(arg_name, action="store_true", help=arg_help) - else: - group.add_argument(arg_name, type=arg_type, default=default, help=arg_help) - - # Add fields from each section with a prefix for context - add_fields_to_parser(f for f in fields(TeacherPolicyCfg) if not is_dataclass(f.type)) - add_fields_to_parser(fields(PolicyCfg)) - add_fields_to_parser(fields(AlgorithmCfg)) - add_fields_to_parser(fields(RunnerCfg)) - - @staticmethod - def from_argparse_args(args: argparse.Namespace) -> "TeacherPolicyCfg": - """Creates an instance from argparse arguments.""" - - def extract_fields(cls: Type) -> dict: - """Helper function to extract fields for a given dataclass type.""" - extracted_fields = { - field.name: getattr(args, TeacherPolicyCfg.args_prefix() + field.name) - for field in fields(cls) - if hasattr(args, TeacherPolicyCfg.args_prefix() + field.name) - } - return extracted_fields - - policy_args = extract_fields(PolicyCfg) - algorithm_args = extract_fields(AlgorithmCfg) - runner_args = extract_fields(RunnerCfg) - - return TeacherPolicyCfg( - seed=getattr(args, f"{TeacherPolicyCfg.args_prefix()}seed"), - policy=PolicyCfg(**policy_args), - algorithm=AlgorithmCfg(**algorithm_args), - runner=RunnerCfg(**runner_args), - ) - - def overwrite_policy_cfg_from_file(self, file_path: str): - try: - saved_config = TeacherPolicyCfg.load(file_path) - self.policy = saved_config.policy - except FileNotFoundError: - pass diff --git a/scripts/rsl_rl/train_student_policy.py b/scripts/rsl_rl/train_student_policy.py index 9b4edda..3aa045f 100644 --- a/scripts/rsl_rl/train_student_policy.py +++ b/scripts/rsl_rl/train_student_policy.py @@ -17,46 +17,56 @@ """Launch Isaac Sim Simulator first.""" import argparse -import json import os -from dataclasses import MISSING, fields from datetime import datetime -from teacher_policy_cfg import TeacherPolicyCfg -from utils import get_ppo_runner_and_checkpoint_path - -from neural_wbc.student_policy import StudentPolicyTrainer, StudentPolicyTrainerCfg, TeacherPolicy - from isaaclab.app import AppLauncher -student_policy_args_prefix = StudentPolicyTrainerCfg.args_prefix() +import cli_args # isort: skip + project_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../")) # add argparse arguments parser = argparse.ArgumentParser( description="Train student policy from a neural WBC teacher policy.", formatter_class=argparse.ArgumentDefaultsHelpFormatter, ) -parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.") -parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment.") -parser.add_argument("--reference_motion_path", type=str, default=None, help="Path to the reference motion dataset.") -parser.add_argument("--output_dir", type=str, default="logs", help="Directory to store the training output.") -parser.add_argument("--robot", type=str, choices=["h1", "gr1"], default="h1", help="Robot used in environment") parser.add_argument( - f"--{student_policy_args_prefix}root_path", + "--reference_motion_path", type=str, - default=os.path.join(project_dir, "logs/student"), - help="Root directory of all student policy training artifacts.", + default=None, + help="Path to the reference motion dataset.", ) parser.add_argument( - f"--{student_policy_args_prefix}append_timestamp", - action=argparse.BooleanOptionalAction, - default=True, - help="If true the current timestamp will be added to the output student policy path.", + "--teacher_load_run", + type=str, + default=".*", + help="Experiment name of the teacher. Default is '.*' which fetches the latest experiment in the directory", ) -# append StudentPolicyTrainerCfg arguments -StudentPolicyTrainerCfg.add_args_to_parser(parser) -# append RSL-RL cli arguments -TeacherPolicyCfg.add_args_to_parser(parser) +parser.add_argument( + "--teacher_load_checkpoint", + type=str, + default="model_.*.pt", + help="Checkpoint of the teacher. Default is 'model_.*.pt' which fetches the latest checkpoint in the run directory", +) +parser.add_argument( + "--output_dir", + type=str, + default="logs", + help="Directory to store the training output.", +) +parser.add_argument( + "--robot", + type=str, + choices=["h1", "gr1"], + default="h1", + help="Robot used in environment", +) +parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.") +parser.add_argument("--max_iterations", type=int, default=None, help="RL Policy training iterations.") +parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment") + +cli_args.add_rsl_rl_args(parser) + # append AppLauncher cli args AppLauncher.add_app_launcher_args(parser) args_cli = parser.parse_args() @@ -69,6 +79,9 @@ import torch +from rsl_rl.runners import OnPolicyRunner +from student_policy_cfg import StudentPolicyRunnerCfg + # Import extensions to set up environment tasks from vecenv_wrapper import RslRlNeuralWBCVecEnvWrapper @@ -76,215 +89,63 @@ from neural_wbc.isaac_lab_wrapper.neural_wbc_env import NeuralWBCEnv from neural_wbc.isaac_lab_wrapper.neural_wbc_env_cfg_h1 import NeuralWBCEnvCfgH1 +from isaaclab_tasks.utils import get_checkpoint_path + torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True torch.backends.cudnn.deterministic = False torch.backends.cudnn.benchmark = False -class NeuralWBCTeacherPolicy(TeacherPolicy): - """Class to load and apply a policy trained with neural WBC.""" - - def __init__(self, env: RslRlNeuralWBCVecEnvWrapper, teacher_policy_cfg: TeacherPolicyCfg): - self.path = teacher_policy_cfg.runner.resume_path - # Here we assume that the path to the teacher policy is {experiment_dir}/{teacher_name}. - self.name = self.path.strip("/").split("/")[-1] - - self._ppo_runner, checkpoint_path = get_ppo_runner_and_checkpoint_path( - teacher_policy_cfg=teacher_policy_cfg, wrapped_env=env, device=env.device - ) - self._ppo_runner.load(checkpoint_path) - print(f"[INFO]: Loaded model checkpoint from: {checkpoint_path}") - - self._actor_critic = self._ppo_runner.alg.actor_critic - self._actor_critic.eval() # switch to evaluation mode (dropout for example) - if env.device is not None: - self._actor_critic.to(env.device) - - def act_rollout(self, observations): - # use act from actor_critic - return self._actor_critic.act(observations) - - def act(self, observations): - return self._actor_critic.act_inference(observations) - - -def resolve_student_policy_path(root_path: str, teacher_policy: NeuralWBCTeacherPolicy, append_timestamp: bool): - """ - Generates a path for storing student policy configurations based on the root path and teacher policy. - - This function appends a timestamp and the teacher policy name to the given root path to create a unique path - for storing student policy configurations. - - Args: - root_path (str): The root directory path where the student policy configurations will be stored. - teacher_policy (NeuralWBCTeacherPolicy): The teacher policy instance whose name will be used in the generated path. - - Returns: - str: The generated absolute path for storing the student policy configurations. - """ - path = os.path.join(os.path.abspath(root_path), teacher_policy.name) - if append_timestamp: - path += "_" + datetime.now().strftime("%y_%m_%d_%H-%M-%S") - - return path - - -def get_config_values_from_argparser(args: argparse.Namespace, teacher_policy: NeuralWBCTeacherPolicy): - """ - Extracts configuration values from command-line arguments based on a predefined prefix and field names. - - This function processes the command-line arguments, filters out those that match a specified prefix, - and returns a dictionary of configuration values that are relevant to the StudentPolicyTrainerCfg dataclass. - - Args: - args (argparse.Namespace): The parsed command-line arguments. - - Returns: - dict: A dictionary containing configuration values extracted from the command-line arguments. - """ - # Get the set of all field names in the StudentPolicyTrainerCfg dataclass - field_names = {field_info.name for field_info in fields(StudentPolicyTrainerCfg)} - - # Extract arguments from the args object and remove the prefix - args_dict = { - key.removeprefix(student_policy_args_prefix): value - for key, value in vars(args).items() - if key.startswith(student_policy_args_prefix) and key.removeprefix(student_policy_args_prefix) in field_names - } - - args_dict["teacher_policy"] = teacher_policy - - # Set student policy path - args_dict["student_policy_path"] = resolve_student_policy_path( - root_path=getattr(args, f"{student_policy_args_prefix}root_path"), - teacher_policy=teacher_policy, - append_timestamp=getattr(args, f"{student_policy_args_prefix}append_timestamp"), - ) - - return args_dict - - -def get_student_trainer_cfg(args: argparse.Namespace, env: NeuralWBCEnv, teacher_policy: NeuralWBCTeacherPolicy): - """ - Create an instance of StudentPolicyTrainerCfg from command-line arguments, environment configuration, and a teacher policy. - - This function processes the command-line arguments, extracts necessary values, fills in any missing values using - the environment configuration, and creates an instance of the StudentPolicyTrainerCfg dataclass. - - Args: - args (argparse.Namespace): The parsed command-line arguments. - env (NeuralWBCEnv): The environment object that provides necessary configuration values. - teacher_policy (NeuralWBCTeacherPolicy): The teacher policy instance to be included in the configuration. - - Returns: - StudentPolicyTrainerCfg: An instance of the StudentPolicyTrainerCfg dataclass. - - Raises: - ValueError: If a required field does not have a default value and is not provided in the arguments or the environment configuration. - """ - # First try loading pre-existing config. - previous_config = load_student_policy_trainer_cfg(args=args, teacher_policy=teacher_policy) - if previous_config: - return previous_config - - args_dict = get_config_values_from_argparser(args=args, teacher_policy=teacher_policy) - - # Identify fields that do not have default values and are not provided in the arguments - no_default_fields = [ - field_info.name - for field_info in fields(StudentPolicyTrainerCfg) - if args_dict.get(field_info.name) is None - and field_info.default is MISSING - and field_info.default_factory is MISSING - ] - - # Fill in missing values for fields that do not have defaults - for name in no_default_fields: - value = None - if name == "num_policy_obs": - # Special case: get the value from the environment's num_observations attribute - value = env.num_observations - elif hasattr(env, name): - # Check if the environment has an attribute matching the field name - value = getattr(env, name) - elif hasattr(env.cfg, name): - # Check if the environment's configuration has an attribute matching the field name - value = getattr(env.cfg, name) - else: - # Raise an error if the field value cannot be determined - raise ValueError( - f"student_policy.{name} does not have a default value in the student training configuration or the" - " environment. Please specify a value." - ) - # Update the args_dict with the determined value - args_dict[name] = value - - return StudentPolicyTrainerCfg(**args_dict) - - -def load_student_policy_trainer_cfg(args: argparse.Namespace, teacher_policy: NeuralWBCTeacherPolicy): - """ - Loads the student policy trainer configuration from a specified resume path. - - This function checks if a resume path is provided in the command-line arguments. If so, it loads the configuration - from a config.json at the resume path, updates it with any new arguments provided, and returns an instance of - StudentPolicyTrainerCfg. - - Args: - args (argparse.Namespace): The parsed command-line arguments. - teacher_policy (NeuralWBCTeacherPolicy): The teacher policy instance to be included in the configuration. - - Returns: - StudentPolicyTrainerCfg or None: An instance of the StudentPolicyTrainerCfg dataclass if a resume path is provided, - otherwise None. - """ - # Check if resume_path is specified by the user. - resume_path = getattr(args, f"{student_policy_args_prefix}resume_path") - if not resume_path: - return None - - with open(os.path.join(resume_path, "config.json")) as fh: - config_dict = json.load(fh) - - del config_dict["teacher_policy"] - - args_dict = get_config_values_from_argparser(args=args, teacher_policy=teacher_policy) - for key, value in args_dict.items(): - # Only overwrite values from arguments that are not None. - if value is not None: - config_dict[key] = value - - return StudentPolicyTrainerCfg(**config_dict) - - def main(): # parse configuration if args_cli.robot == "h1": env_cfg = NeuralWBCEnvCfgH1(mode=NeuralWBCModes.DISTILL) elif args_cli.robot == "gr1": raise ValueError("GR1 is not yet implemented") + agent_cfg = cli_args.update_rsl_rl_cfg(StudentPolicyRunnerCfg(), args_cli) + env_cfg.scene.num_envs = args_cli.num_envs env_cfg.scene.env_spacing = 20 env_cfg.terrain.env_spacing = 20 if args_cli.reference_motion_path: env_cfg.reference_motion_manager.motion_path = args_cli.reference_motion_path + env_cfg.seed = agent_cfg.seed + env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device + # Create env and wrap it for RSL RL. env = NeuralWBCEnv(cfg=env_cfg) - wrapped_env = RslRlNeuralWBCVecEnvWrapper(env) + env = RslRlNeuralWBCVecEnvWrapper(env) + + # get the log path of the teacher policy + log_root_path = os.path.join("logs", "rsl_rl", args_cli.robot, agent_cfg.experiment_name) + log_root_path = os.path.abspath(log_root_path) + teacher_root_path = os.path.join("logs", "rsl_rl", args_cli.robot, "teacher_policy") + teacher_root_path = os.path.abspath(teacher_root_path) + print(f"[INFO] Logging experiment in directory: {log_root_path}") + print(f"[INFO] Loading teacher policy from: {teacher_root_path}") + # specify directory for logging runs: {time-stamp}_{run_name} + log_dir = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + log_dir = os.path.join(log_root_path, log_dir) + + trained_teacher_path = get_checkpoint_path( + teacher_root_path, args_cli.teacher_load_run, args_cli.teacher_load_checkpoint + ) - teacher_policy_cfg = TeacherPolicyCfg.from_argparse_args(args_cli) - teacher_policy = NeuralWBCTeacherPolicy(env=wrapped_env, teacher_policy_cfg=teacher_policy_cfg) + if agent_cfg.resume: + resume_path = get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint) # set seed of the environment - env.seed(teacher_policy_cfg.seed) + env.seed(agent_cfg.seed) + + distillation_runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=log_dir, device=env.unwrapped.device) + + distillation_runner.load(trained_teacher_path) + if agent_cfg.resume: + distillation_runner.load(resume_path) - trainer_cfg = get_student_trainer_cfg(args=args_cli, env=env, teacher_policy=teacher_policy) - os.makedirs(trainer_cfg.student_policy_path, exist_ok=True) - env_cfg.save(os.path.join(trainer_cfg.student_policy_path, "env_config.json")) - trainer = StudentPolicyTrainer(env=wrapped_env, cfg=trainer_cfg) - trainer.run() + distillation_runner.learn(num_learning_iterations=agent_cfg.max_iterations, init_at_random_ep_len=True) # close the simulator env.close() diff --git a/scripts/rsl_rl/train_teacher_policy.py b/scripts/rsl_rl/train_teacher_policy.py index 641320f..707b427 100644 --- a/scripts/rsl_rl/train_teacher_policy.py +++ b/scripts/rsl_rl/train_teacher_policy.py @@ -14,31 +14,39 @@ # limitations under the License. -"""Script to train RL agent with RSL-RL.""" +"""Script to train RL teacher with RSL-RL.""" """Launch Isaac Sim Simulator first.""" import argparse +from isaaclab.app import AppLauncher + # local imports -from teacher_policy_cfg import TeacherPolicyCfg +import cli_args # isort: skip -from isaaclab.app import AppLauncher # add argparse arguments -parser = argparse.ArgumentParser( - description="Train an RL agent with RSL-RL.", formatter_class=argparse.ArgumentDefaultsHelpFormatter -) +parser = argparse.ArgumentParser(description="Train an RL agent with RSL-RL.") +parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.") +parser.add_argument("--video_length", type=int, default=200, help="Length of the recorded video (in steps).") +parser.add_argument("--video_interval", type=int, default=2000, help="Interval between video recordings (in steps).") parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.") +parser.add_argument("--task", type=str, default=None, help="Name of the task.") parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment") +parser.add_argument("--max_iterations", type=int, default=None, help="RL Policy training iterations.") +parser.add_argument( + "--distributed", action="store_true", default=False, help="Run training with multiple GPUs or nodes." +) parser.add_argument("--reference_motion_path", type=str, default=None, help="Path to the reference motion dataset.") parser.add_argument("--robot", type=str, choices=["h1", "gr1"], default="h1", help="Robot used in environment") # append RSL-RL cli arguments -TeacherPolicyCfg.add_args_to_parser(parser) +cli_args.add_rsl_rl_args(parser) + # append AppLauncher cli args AppLauncher.add_app_launcher_args(parser) -args_cli = parser.parse_args() +args_cli, hydra_args = parser.parse_known_args() # launch omniverse app app_launcher = AppLauncher(args_cli) @@ -46,15 +54,13 @@ """Rest everything follows.""" -from utils import get_customized_rsl_rl, get_ppo_runner_and_checkpoint_path - import os import torch from datetime import datetime # Import your specific module/class -get_customized_rsl_rl() from rsl_rl.runners.on_policy_runner import OnPolicyRunner +from teacher_policy_cfg import TeacherPolicyRunnerCfg # Import extensions to set up environment tasks from vecenv_wrapper import RslRlNeuralWBCVecEnvWrapper @@ -62,6 +68,8 @@ from neural_wbc.isaac_lab_wrapper.neural_wbc_env import NeuralWBCEnv from neural_wbc.isaac_lab_wrapper.neural_wbc_env_cfg_h1 import NeuralWBCEnvCfgH1 +from isaaclab_tasks.utils import get_checkpoint_path + torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True torch.backends.cudnn.deterministic = False @@ -82,37 +90,47 @@ def main(): if args_cli.reference_motion_path: env_cfg.reference_motion_manager.motion_path = args_cli.reference_motion_path - teacher_policy_cfg = TeacherPolicyCfg.from_argparse_args(args_cli) + agent_cfg = cli_args.update_rsl_rl_cfg(TeacherPolicyRunnerCfg(), args_cli) + + env_cfg.seed = agent_cfg.seed + env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device + + # multi-gpu training configuration + if args_cli.distributed: + env_cfg.sim.device = f"cuda:{app_launcher.local_rank}" + agent_cfg.device = f"cuda:{app_launcher.local_rank}" + + # set seed to have diversity in different threads + seed = agent_cfg.seed + app_launcher.local_rank + env_cfg.seed = seed + agent_cfg.seed = seed # Create env and wrap it for RSL RL. env = NeuralWBCEnv(cfg=env_cfg) env = RslRlNeuralWBCVecEnvWrapper(env) - # specify directory for logging experiments - log_dir = os.path.join( - os.path.abspath(teacher_policy_cfg.runner.path), datetime.now().strftime("%y_%m_%d_%H-%M-%S") - ) - os.makedirs(log_dir, exist_ok=True) - print(f"[INFO] Log directory: {log_dir}") - - if teacher_policy_cfg.runner.resume: - ppo_runner, checkpoint_path = get_ppo_runner_and_checkpoint_path( - teacher_policy_cfg=teacher_policy_cfg, wrapped_env=env, device=env.unwrapped.device, log_dir=log_dir - ) - ppo_runner.load(checkpoint_path) - print(f"[INFO]: Loaded model checkpoint from: {checkpoint_path}") - else: - ppo_runner = OnPolicyRunner(env, teacher_policy_cfg.to_dict(), log_dir=log_dir, device=env.unwrapped.device) - - # Store the configuration - teacher_policy_cfg.save(os.path.join(log_dir, "config.json")) - env_cfg.save(os.path.join(log_dir, "env_config.json")) - - # set seed of the environment - env.seed(teacher_policy_cfg.seed) + # specify directory for logging experiments as in isaaclab + log_root_path = os.path.join("logs", "rsl_rl", args_cli.robot, agent_cfg.experiment_name) + log_root_path = os.path.abspath(log_root_path) + print(f"[INFO] Logging experiment in directory: {log_root_path}") + + # specify directory for logging runs: {time-stamp}_{run_name} + log_dir = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + log_dir = os.path.join(log_root_path, log_dir) + + # save resume path before creating a new log_dir + if agent_cfg.resume: + resume_path = get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint) + + ppo_runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=log_dir, device=env.unwrapped.device) + + if agent_cfg.resume: + # load previously trained model + ppo_runner.load(resume_path) + print(f"[INFO]: Loaded model checkpoint from: {resume_path}") # run training - ppo_runner.learn(num_learning_iterations=teacher_policy_cfg.runner.max_iterations, init_at_random_ep_len=True) + ppo_runner.learn(num_learning_iterations=agent_cfg.max_iterations, init_at_random_ep_len=True) # close the simulator env.close() diff --git a/scripts/rsl_rl/utils.py b/scripts/rsl_rl/utils.py index 6942426..65e953d 100644 --- a/scripts/rsl_rl/utils.py +++ b/scripts/rsl_rl/utils.py @@ -39,7 +39,7 @@ def get_customized_rsl_rl(): import pkg_resources - dist = pkg_resources.require("rsl_rl")[0] + dist = pkg_resources.require("rsl-rl-lib")[0] sys.path.insert(0, dist.location) # Remove 'rsl_rl' from sys.modules if it was already imported diff --git a/scripts/rsl_rl/vecenv_wrapper.py b/scripts/rsl_rl/vecenv_wrapper.py index 9516eda..976f780 100644 --- a/scripts/rsl_rl/vecenv_wrapper.py +++ b/scripts/rsl_rl/vecenv_wrapper.py @@ -76,7 +76,7 @@ def __init__(self, env: NeuralWBCEnv): else: self.num_actions = self._env.action_space.shape[1] if hasattr(self._env, "observation_manager"): - self.num_obs = self._env.observation_manager.group_obs_dim["teacher_policy"][0] + self.num_obs = self._env.observation_manager.group_obs_dim["teacher"][0] else: self.num_obs = self._env.num_observations # -- privileged observations @@ -144,12 +144,12 @@ def get_full_observations(self): else: obs_dict = self.unwrapped._get_observations() if self.unwrapped.cfg.observation_noise_model: - obs_dict["teacher_policy"] = self.unwrapped._observation_noise_model.apply(obs_dict["teacher_policy"]) + obs_dict["teacher"] = self.unwrapped._observation_noise_model.apply(obs_dict["teacher"]) return obs_dict def get_teacher_observations(self) -> torch.Tensor: """Returns the current observations of the environment.""" - return self.get_full_observations()["teacher_policy"] + return self.get_full_observations()["teacher"] def get_privileged_observations(self) -> torch.Tensor: """Returns the current observations of the environment.""" @@ -185,7 +185,7 @@ def reset(self) -> tuple[torch.Tensor, torch.Tensor]: # noqa: D102 # return observations if self._mode.is_distill_mode(): return obs_dict["student_policy"], torch.tensor([]) - return obs_dict["teacher_policy"], obs_dict["critic"] + return obs_dict["teacher"], obs_dict["critic"] def step(self, actions: torch.Tensor): # record step information @@ -193,14 +193,13 @@ def step(self, actions: torch.Tensor): # compute dones for compatibility with RSL-RL dones = (terminated | truncated).to(dtype=torch.long) # move extra observations to the extras dict - obs = obs_dict["student_policy"] if self._mode.is_distill_mode() else obs_dict["teacher_policy"] - privileged_obs = obs_dict["critic"] + obs = obs_dict["student_policy"] if self._mode.is_distill_mode() else obs_dict["teacher"] extras["observations"] = obs_dict # move time out information to the extras dict extras["time_outs"] = truncated # return the step information - return obs, privileged_obs, rew, dones, extras + return obs, rew, dones, extras def close(self): # noqa: D102 return self._env.close() diff --git a/third_party/rsl_rl/.gitignore b/third_party/rsl_rl/.gitignore deleted file mode 100755 index 34f42f1..0000000 --- a/third_party/rsl_rl/.gitignore +++ /dev/null @@ -1,12 +0,0 @@ -# IDEs -.idea - -# builds -*.egg-info - -# cache -__pycache__ -.pytest_cache - -# vs code -.vscode \ No newline at end of file diff --git a/third_party/rsl_rl/LICENSE b/third_party/rsl_rl/LICENSE deleted file mode 100755 index 01d9567..0000000 --- a/third_party/rsl_rl/LICENSE +++ /dev/null @@ -1,30 +0,0 @@ -Copyright (c) 2021, ETH Zurich, Nikita Rudin -Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES -All rights reserved. - -Redistribution and use in source and binary forms, with or without modification, -are permitted provided that the following conditions are met: - -1. Redistributions of source code must retain the above copyright notice, - this list of conditions and the following disclaimer. - -2. Redistributions in binary form must reproduce the above copyright notice, - this list of conditions and the following disclaimer in the documentation - and/or other materials provided with the distribution. - -3. Neither the name of the copyright holder nor the names of its contributors - may be used to endorse or promote products derived from this software without - specific prior written permission. - -THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND -ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED -WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR -ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES -(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; -LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND -ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT -(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS -SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - -See licenses/dependencies for license information of dependencies of this package. \ No newline at end of file diff --git a/third_party/rsl_rl/README.md b/third_party/rsl_rl/README.md deleted file mode 100755 index c4a5bcf..0000000 --- a/third_party/rsl_rl/README.md +++ /dev/null @@ -1,23 +0,0 @@ -# RSL RL -Fast and simple implementation of RL algorithms, designed to run fully on GPU. -This code is an evolution of `rl-pytorch` provided with NVIDIA's Isaac GYM. - -Only PPO is implemented for now. More algorithms will be added later. -Contributions are welcome. - -## Setup - -``` -git clone https://github.com/leggedrobotics/rsl_rl -cd rsl_rl -pip install -e . -``` - -### Useful Links ### -Example use case: https://github.com/leggedrobotics/legged_gym -Project website: https://leggedrobotics.github.io/legged_gym/ -Paper: https://arxiv.org/abs/2109.11978 - -**Maintainer**: Nikita Rudin -**Affiliation**: Robotic Systems Lab, ETH Zurich & NVIDIA -**Contact**: rudinn@ethz.ch \ No newline at end of file diff --git a/third_party/rsl_rl/licenses/dependencies/numpy_license.txt b/third_party/rsl_rl/licenses/dependencies/numpy_license.txt deleted file mode 100755 index 84e9bfe..0000000 --- a/third_party/rsl_rl/licenses/dependencies/numpy_license.txt +++ /dev/null @@ -1,30 +0,0 @@ -Copyright (c) 2005-2021, NumPy Developers. -All rights reserved. - -Redistribution and use in source and binary forms, with or without -modification, are permitted provided that the following conditions are -met: - - * Redistributions of source code must retain the above copyright - notice, this list of conditions and the following disclaimer. - - * Redistributions in binary form must reproduce the above - copyright notice, this list of conditions and the following - disclaimer in the documentation and/or other materials provided - with the distribution. - - * Neither the name of the NumPy Developers nor the names of any - contributors may be used to endorse or promote products derived - from this software without specific prior written permission. - -THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS -"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT -LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR -A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT -OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, -SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT -LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, -DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY -THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT -(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. \ No newline at end of file diff --git a/third_party/rsl_rl/licenses/dependencies/torch_license.txt b/third_party/rsl_rl/licenses/dependencies/torch_license.txt deleted file mode 100755 index 244b249..0000000 --- a/third_party/rsl_rl/licenses/dependencies/torch_license.txt +++ /dev/null @@ -1,73 +0,0 @@ -From PyTorch: - -Copyright (c) 2016- Facebook, Inc (Adam Paszke) -Copyright (c) 2014- Facebook, Inc (Soumith Chintala) -Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert) -Copyright (c) 2012-2014 Deepmind Technologies (Koray Kavukcuoglu) -Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu) -Copyright (c) 2011-2013 NYU (Clement Farabet) -Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston) -Copyright (c) 2006 Idiap Research Institute (Samy Bengio) -Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz) - -From Caffe2: - -Copyright (c) 2016-present, Facebook Inc. All rights reserved. - -All contributions by Facebook: -Copyright (c) 2016 Facebook Inc. - -All contributions by Google: -Copyright (c) 2015 Google Inc. -All rights reserved. - -All contributions by Yangqing Jia: -Copyright (c) 2015 Yangqing Jia -All rights reserved. - -All contributions by Kakao Brain: -Copyright 2019-2020 Kakao Brain - -All contributions from Caffe: -Copyright(c) 2013, 2014, 2015, the respective contributors -All rights reserved. - -All other contributions: -Copyright(c) 2015, 2016 the respective contributors -All rights reserved. - -Caffe2 uses a copyright model similar to Caffe: each contributor holds -copyright over their contributions to Caffe2. The project versioning records -all such contribution and copyright details. If a contributor wants to further -mark their specific copyright on a particular contribution, they should -indicate their copyright solely in the commit message of the change when it is -committed. - -All rights reserved. - -Redistribution and use in source and binary forms, with or without -modification, are permitted provided that the following conditions are met: - -1. Redistributions of source code must retain the above copyright - notice, this list of conditions and the following disclaimer. - -2. Redistributions in binary form must reproduce the above copyright - notice, this list of conditions and the following disclaimer in the - documentation and/or other materials provided with the distribution. - -3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America - and IDIAP Research Institute nor the names of its contributors may be - used to endorse or promote products derived from this software without - specific prior written permission. - -THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE -ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE -LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR -CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF -SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS -INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN -CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) -ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE -POSSIBILITY OF SUCH DAMAGE. \ No newline at end of file diff --git a/third_party/rsl_rl/rsl_rl/__init__.py b/third_party/rsl_rl/rsl_rl/__init__.py deleted file mode 100755 index 466dfca..0000000 --- a/third_party/rsl_rl/rsl_rl/__init__.py +++ /dev/null @@ -1,29 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: BSD-3-Clause -# -# Redistribution and use in source and binary forms, with or without -# modification, are permitted provided that the following conditions are met: -# -# 1. Redistributions of source code must retain the above copyright notice, this -# list of conditions and the following disclaimer. -# -# 2. Redistributions in binary form must reproduce the above copyright notice, -# this list of conditions and the following disclaimer in the documentation -# and/or other materials provided with the distribution. -# -# 3. Neither the name of the copyright holder nor the names of its -# contributors may be used to endorse or promote products derived from -# this software without specific prior written permission. -# -# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE -# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR -# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER -# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, -# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -# -# Copyright (c) 2021 ETH Zurich, Nikita Rudin \ No newline at end of file diff --git a/third_party/rsl_rl/rsl_rl/algorithms/__init__.py b/third_party/rsl_rl/rsl_rl/algorithms/__init__.py deleted file mode 100755 index 93542af..0000000 --- a/third_party/rsl_rl/rsl_rl/algorithms/__init__.py +++ /dev/null @@ -1,31 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: BSD-3-Clause -# -# Redistribution and use in source and binary forms, with or without -# modification, are permitted provided that the following conditions are met: -# -# 1. Redistributions of source code must retain the above copyright notice, this -# list of conditions and the following disclaimer. -# -# 2. Redistributions in binary form must reproduce the above copyright notice, -# this list of conditions and the following disclaimer in the documentation -# and/or other materials provided with the distribution. -# -# 3. Neither the name of the copyright holder nor the names of its -# contributors may be used to endorse or promote products derived from -# this software without specific prior written permission. -# -# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE -# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR -# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER -# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, -# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -# -# Copyright (c) 2021 ETH Zurich, Nikita Rudin - -from .ppo import PPO diff --git a/third_party/rsl_rl/rsl_rl/algorithms/ppo.py b/third_party/rsl_rl/rsl_rl/algorithms/ppo.py deleted file mode 100755 index c2fc9e1..0000000 --- a/third_party/rsl_rl/rsl_rl/algorithms/ppo.py +++ /dev/null @@ -1,220 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: BSD-3-Clause -# -# Redistribution and use in source and binary forms, with or without -# modification, are permitted provided that the following conditions are met: -# -# 1. Redistributions of source code must retain the above copyright notice, this -# list of conditions and the following disclaimer. -# -# 2. Redistributions in binary form must reproduce the above copyright notice, -# this list of conditions and the following disclaimer in the documentation -# and/or other materials provided with the distribution. -# -# 3. Neither the name of the copyright holder nor the names of its -# contributors may be used to endorse or promote products derived from -# this software without specific prior written permission. -# -# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE -# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR -# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER -# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, -# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -# -# Copyright (c) 2021 ETH Zurich, Nikita Rudin - -import torch -import torch.nn as nn -import torch.optim as optim - -from rsl_rl.modules import ActorCritic -from rsl_rl.storage import RolloutStorage - - -class PPO: - actor_critic: ActorCritic - - def __init__( - self, - actor_critic, - num_learning_epochs=1, - num_mini_batches=1, - clip_param=0.2, - gamma=0.998, - lam=0.95, - value_loss_coef=1.0, - entropy_coef=0.0, - learning_rate=1e-3, - max_grad_norm=1.0, - use_clipped_value_loss=True, - schedule="fixed", - desired_kl=0.01, - device="cpu", - **kwargs - ): - - self.device = device - - self.desired_kl = desired_kl - self.schedule = schedule - self.learning_rate = learning_rate - - # PPO components - self.actor_critic = actor_critic - self.actor_critic.to(self.device) - self.storage = None # initialized later - self.optimizer = optim.Adam(self.actor_critic.parameters(), lr=learning_rate) - self.transition = RolloutStorage.Transition() - - # PPO parameters - self.clip_param = clip_param - self.num_learning_epochs = num_learning_epochs - self.num_mini_batches = num_mini_batches - self.value_loss_coef = value_loss_coef - self.entropy_coef = entropy_coef - self.gamma = gamma - self.lam = lam - self.max_grad_norm = max_grad_norm - self.use_clipped_value_loss = use_clipped_value_loss - - def init_storage(self, num_envs, num_transitions_per_env, actor_obs_shape, critic_obs_shape, action_shape): - - self.storage = RolloutStorage( - num_envs, - num_transitions_per_env, - actor_obs_shape, - critic_obs_shape, - action_shape, - device=self.device, - ) - - def test_mode(self): - self.actor_critic.test() - - def train_mode(self): - self.actor_critic.train() - - def act(self, obs, critic_obs): - if self.actor_critic.is_recurrent: - self.transition.hidden_states = self.actor_critic.get_hidden_states() - # Compute the actions and values - self.transition.actions = self.actor_critic.act(obs).detach() - self.transition.values = self.actor_critic.evaluate(critic_obs).detach() - self.transition.actions_log_prob = self.actor_critic.get_actions_log_prob(self.transition.actions).detach() - self.transition.action_mean = self.actor_critic.action_mean.detach() - self.transition.action_sigma = self.actor_critic.action_std.detach() - # need to record obs and critic_obs before env.step() - self.transition.observations = obs - self.transition.critic_observations = critic_obs - return self.transition.actions - - def process_env_step(self, rewards, dones, infos): - self.transition.rewards = rewards.clone() - self.transition.dones = dones - # Bootstrapping on time outs - if "time_outs" in infos: - self.transition.rewards += self.gamma * torch.squeeze( - self.transition.values * infos["time_outs"].unsqueeze(1).to(self.device), 1 - ) - - # Record the transition - self.storage.add_transitions(self.transition) - self.transition.clear() - self.actor_critic.reset(dones) - - def compute_returns(self, last_critic_obs): - last_values = self.actor_critic.evaluate(last_critic_obs).detach() - self.storage.compute_returns(last_values, self.gamma, self.lam) - - def update(self, epoch=0): - mean_value_loss = 0 - mean_surrogate_loss = 0 - - if self.actor_critic.is_recurrent: - generator = self.storage.reccurent_mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) - else: - generator = self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) - - for ( - obs_batch, - critic_obs_batch, - actions_batch, - target_values_batch, - advantages_batch, - returns_batch, - old_actions_log_prob_batch, - old_mu_batch, - old_sigma_batch, - hid_states_batch, - masks_batch, - obs_squential, - kin_dict_batch, - ) in generator: - - self.actor_critic.act(obs_batch) - actions_log_prob_batch = self.actor_critic.get_actions_log_prob(actions_batch) - value_batch = self.actor_critic.evaluate(critic_obs_batch) - mu_batch = self.actor_critic.action_mean - sigma_batch = self.actor_critic.action_std - entropy_batch = self.actor_critic.entropy - - # KL - if self.desired_kl is not None and self.schedule == "adaptive": - with torch.inference_mode(): - kl = torch.sum( - torch.log(sigma_batch / old_sigma_batch + 1.0e-5) - + (torch.square(old_sigma_batch) + torch.square(old_mu_batch - mu_batch)) - / (2.0 * torch.square(sigma_batch)) - - 0.5, - axis=-1, - ) - kl_mean = torch.mean(kl) - - if kl_mean > self.desired_kl * 2.0: - self.learning_rate = max(1e-5, self.learning_rate / 1.5) - elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0: - self.learning_rate = min(1e-2, self.learning_rate * 1.5) - - for param_group in self.optimizer.param_groups: - param_group["lr"] = self.learning_rate - - # Surrogate loss - ratio = torch.exp(actions_log_prob_batch - torch.squeeze(old_actions_log_prob_batch)) - surrogate = -torch.squeeze(advantages_batch) * ratio - surrogate_clipped = -torch.squeeze(advantages_batch) * torch.clamp( - ratio, 1.0 - self.clip_param, 1.0 + self.clip_param - ) - surrogate_loss = torch.max(surrogate, surrogate_clipped).mean() - - # Value function loss - if self.use_clipped_value_loss: - value_clipped = target_values_batch + (value_batch - target_values_batch).clamp( - -self.clip_param, self.clip_param - ) - value_losses = (value_batch - returns_batch).pow(2) - value_losses_clipped = (value_clipped - returns_batch).pow(2) - value_loss = torch.max(value_losses, value_losses_clipped).mean() - else: - value_loss = (returns_batch - value_batch).pow(2).mean() - - loss = surrogate_loss + self.value_loss_coef * value_loss - self.entropy_coef * entropy_batch.mean() - # Gradient step - self.optimizer.zero_grad() - loss.backward() - nn.utils.clip_grad_norm_(self.actor_critic.parameters(), self.max_grad_norm) - self.optimizer.step() - - mean_value_loss += value_loss.item() - mean_surrogate_loss += surrogate_loss.item() - - num_updates = self.num_learning_epochs * self.num_mini_batches - mean_value_loss /= num_updates - mean_surrogate_loss /= num_updates - - self.storage.clear() - return mean_value_loss, mean_surrogate_loss diff --git a/third_party/rsl_rl/rsl_rl/env/__init__.py b/third_party/rsl_rl/rsl_rl/env/__init__.py deleted file mode 100755 index 9539b9f..0000000 --- a/third_party/rsl_rl/rsl_rl/env/__init__.py +++ /dev/null @@ -1,31 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: BSD-3-Clause -# -# Redistribution and use in source and binary forms, with or without -# modification, are permitted provided that the following conditions are met: -# -# 1. Redistributions of source code must retain the above copyright notice, this -# list of conditions and the following disclaimer. -# -# 2. Redistributions in binary form must reproduce the above copyright notice, -# this list of conditions and the following disclaimer in the documentation -# and/or other materials provided with the distribution. -# -# 3. Neither the name of the copyright holder nor the names of its -# contributors may be used to endorse or promote products derived from -# this software without specific prior written permission. -# -# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE -# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR -# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER -# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, -# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -# -# Copyright (c) 2021 ETH Zurich, Nikita Rudin - -from .vec_env import VecEnv \ No newline at end of file diff --git a/third_party/rsl_rl/rsl_rl/env/vec_env.py b/third_party/rsl_rl/rsl_rl/env/vec_env.py deleted file mode 100755 index 6a7ef1b..0000000 --- a/third_party/rsl_rl/rsl_rl/env/vec_env.py +++ /dev/null @@ -1,60 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: BSD-3-Clause -# -# Redistribution and use in source and binary forms, with or without -# modification, are permitted provided that the following conditions are met: -# -# 1. Redistributions of source code must retain the above copyright notice, this -# list of conditions and the following disclaimer. -# -# 2. Redistributions in binary form must reproduce the above copyright notice, -# this list of conditions and the following disclaimer in the documentation -# and/or other materials provided with the distribution. -# -# 3. Neither the name of the copyright holder nor the names of its -# contributors may be used to endorse or promote products derived from -# this software without specific prior written permission. -# -# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE -# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR -# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER -# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, -# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -# -# Copyright (c) 2021 ETH Zurich, Nikita Rudin - -from abc import ABC, abstractmethod -import torch -from typing import Tuple, Union - -# minimal interface of the environment -class VecEnv(ABC): - num_envs: int - num_obs: int - num_privileged_obs: int - num_actions: int - max_episode_length: int - privileged_obs_buf: torch.Tensor - obs_buf: torch.Tensor - rew_buf: torch.Tensor - reset_buf: torch.Tensor - episode_length_buf: torch.Tensor # current episode duration - extras: dict - device: torch.device - @abstractmethod - def step(self, actions: torch.Tensor) -> Tuple[torch.Tensor, Union[torch.Tensor, None], torch.Tensor, torch.Tensor, dict]: - pass - @abstractmethod - def reset(self, env_ids: Union[list, torch.Tensor]): - pass - @abstractmethod - def get_observations(self) -> torch.Tensor: - pass - @abstractmethod - def get_privileged_observations(self) -> Union[torch.Tensor, None]: - pass \ No newline at end of file diff --git a/third_party/rsl_rl/rsl_rl/modules/__init__.py b/third_party/rsl_rl/rsl_rl/modules/__init__.py deleted file mode 100755 index 9befecc..0000000 --- a/third_party/rsl_rl/rsl_rl/modules/__init__.py +++ /dev/null @@ -1,31 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: BSD-3-Clause -# -# Redistribution and use in source and binary forms, with or without -# modification, are permitted provided that the following conditions are met: -# -# 1. Redistributions of source code must retain the above copyright notice, this -# list of conditions and the following disclaimer. -# -# 2. Redistributions in binary form must reproduce the above copyright notice, -# this list of conditions and the following disclaimer in the documentation -# and/or other materials provided with the distribution. -# -# 3. Neither the name of the copyright holder nor the names of its -# contributors may be used to endorse or promote products derived from -# this software without specific prior written permission. -# -# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE -# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR -# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER -# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, -# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -# -# Copyright (c) 2021 ETH Zurich, Nikita Rudin - -from .actor_critic import ActorCritic diff --git a/third_party/rsl_rl/rsl_rl/modules/actor_critic.py b/third_party/rsl_rl/rsl_rl/modules/actor_critic.py deleted file mode 100755 index 007a7c8..0000000 --- a/third_party/rsl_rl/rsl_rl/modules/actor_critic.py +++ /dev/null @@ -1,164 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: BSD-3-Clause -# -# Redistribution and use in source and binary forms, with or without -# modification, are permitted provided that the following conditions are met: -# -# 1. Redistributions of source code must retain the above copyright notice, this -# list of conditions and the following disclaimer. -# -# 2. Redistributions in binary form must reproduce the above copyright notice, -# this list of conditions and the following disclaimer in the documentation -# and/or other materials provided with the distribution. -# -# 3. Neither the name of the copyright holder nor the names of its -# contributors may be used to endorse or promote products derived from -# this software without specific prior written permission. -# -# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE -# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR -# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER -# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, -# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -# -# Copyright (c) 2021 ETH Zurich, Nikita Rudin - - -import torch -import torch.nn as nn -from torch.distributions import Normal - - -class ActorCritic(nn.Module): - is_recurrent = False - - def __init__( - self, - num_actor_obs, - num_critic_obs, - num_actions, - actor_hidden_dims=[256, 256, 256], - critic_hidden_dims=[256, 256, 256], - activation="elu", - init_noise_std=1.0, - **kwargs, - ): - if kwargs: - print( - "ActorCritic.__init__ got unexpected arguments, which will be ignored: " - + str([key for key in kwargs.keys()]) - ) - super(ActorCritic, self).__init__() - - activation = get_activation(activation) - - mlp_input_dim_a = num_actor_obs - mlp_input_dim_c = num_critic_obs - - # Policy - actor_layers = [] - actor_layers.append(nn.Linear(mlp_input_dim_a, actor_hidden_dims[0])) - actor_layers.append(activation) - for l in range(len(actor_hidden_dims)): - if l == len(actor_hidden_dims) - 1: - actor_layers.append(nn.Linear(actor_hidden_dims[l], num_actions)) - else: - actor_layers.append(nn.Linear(actor_hidden_dims[l], actor_hidden_dims[l + 1])) - actor_layers.append(activation) - self.actor = nn.Sequential(*actor_layers) - - # Value function - critic_layers = [] - critic_layers.append(nn.Linear(mlp_input_dim_c, critic_hidden_dims[0])) - critic_layers.append(activation) - for l in range(len(critic_hidden_dims)): - if l == len(critic_hidden_dims) - 1: - critic_layers.append(nn.Linear(critic_hidden_dims[l], 1)) - else: - critic_layers.append(nn.Linear(critic_hidden_dims[l], critic_hidden_dims[l + 1])) - critic_layers.append(activation) - self.critic = nn.Sequential(*critic_layers) - - print(f"Actor MLP: {self.actor}") - print(f"Critic MLP: {self.critic}") - - # Action noise - self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) - self.std.requires_grad = True - self.distribution = None - # disable args validation for speedup - Normal.set_default_validate_args = False - - # seems that we get better performance without init - # self.init_memory_weights(self.memory_a, 0.001, 0.) - # self.init_memory_weights(self.memory_c, 0.001, 0.) - - @staticmethod - # not used at the moment - def init_weights(sequential, scales): - [ - torch.nn.init.orthogonal_(module.weight, gain=scales[idx]) - for idx, module in enumerate(mod for mod in sequential if isinstance(mod, nn.Linear)) - ] - - def reset(self, dones=None): - pass - - def forward(self): - raise NotImplementedError - - @property - def action_mean(self): - return self.distribution.mean - - @property - def action_std(self): - return self.distribution.stddev - - @property - def entropy(self): - return self.distribution.entropy().sum(dim=-1) - - def update_distribution(self, observations): - mean = self.actor(observations) - self.distribution = Normal(mean, mean * 0.0 + self.std) - - def act(self, observations, **kwargs): - self.update_distribution(observations) - return self.distribution.sample() - - def get_actions_log_prob(self, actions): - return self.distribution.log_prob(actions).sum(dim=-1) - - def act_inference(self, observations): - actions_mean = self.actor(observations) - return actions_mean - - def evaluate(self, critic_observations, **kwargs): - value = self.critic(critic_observations) - return value - - -def get_activation(act_name): - if act_name == "elu": - return nn.ELU() - elif act_name == "selu": - return nn.SELU() - elif act_name == "relu": - return nn.ReLU() - elif act_name == "crelu": - return nn.ReLU() - elif act_name == "lrelu": - return nn.LeakyReLU() - elif act_name == "tanh": - return nn.Tanh() - elif act_name == "sigmoid": - return nn.Sigmoid() - else: - print("invalid activation function!") - return None diff --git a/third_party/rsl_rl/rsl_rl/runners/__init__.py b/third_party/rsl_rl/rsl_rl/runners/__init__.py deleted file mode 100755 index 8645ac9..0000000 --- a/third_party/rsl_rl/rsl_rl/runners/__init__.py +++ /dev/null @@ -1,31 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: BSD-3-Clause -# -# Redistribution and use in source and binary forms, with or without -# modification, are permitted provided that the following conditions are met: -# -# 1. Redistributions of source code must retain the above copyright notice, this -# list of conditions and the following disclaimer. -# -# 2. Redistributions in binary form must reproduce the above copyright notice, -# this list of conditions and the following disclaimer in the documentation -# and/or other materials provided with the distribution. -# -# 3. Neither the name of the copyright holder nor the names of its -# contributors may be used to endorse or promote products derived from -# this software without specific prior written permission. -# -# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE -# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR -# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER -# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, -# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -# -# Copyright (c) 2021 ETH Zurich, Nikita Rudin - -from .on_policy_runner import OnPolicyRunner diff --git a/third_party/rsl_rl/rsl_rl/runners/on_policy_runner.py b/third_party/rsl_rl/rsl_rl/runners/on_policy_runner.py deleted file mode 100755 index 2fcaa5d..0000000 --- a/third_party/rsl_rl/rsl_rl/runners/on_policy_runner.py +++ /dev/null @@ -1,302 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: BSD-3-Clause -# -# Redistribution and use in source and binary forms, with or without -# modification, are permitted provided that the following conditions are met: -# -# 1. Redistributions of source code must retain the above copyright notice, this -# list of conditions and the following disclaimer. -# -# 2. Redistributions in binary form must reproduce the above copyright notice, -# this list of conditions and the following disclaimer in the documentation -# and/or other materials provided with the distribution. -# -# 3. Neither the name of the copyright holder nor the names of its -# contributors may be used to endorse or promote products derived from -# this software without specific prior written permission. -# -# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE -# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR -# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER -# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, -# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -# -# Copyright (c) 2021 ETH Zurich, Nikita Rudin - -import time -import os -from collections import deque -import statistics - -import os -import sys - -sys.path.append(os.getcwd()) - -from torch.utils.tensorboard import SummaryWriter -import torch - -from rsl_rl.algorithms import PPO -from rsl_rl.modules import ActorCritic -from rsl_rl.env import VecEnv - - -class OnPolicyRunner: - - def __init__(self, env: VecEnv, train_cfg, log_dir=None, device="cpu"): - - self.cfg = train_cfg["runner"] - self.alg_cfg = train_cfg["algorithm"] - self.policy_cfg = train_cfg["policy"] - - self.device = device - self.env = env - if self.env.num_privileged_obs is not None: - num_critic_obs = self.env.num_privileged_obs - else: - num_critic_obs = self.env.num_obs - actor_critic_class = eval(self.cfg["policy_class_name"]) # ActorCritic - - actor_critic: ActorCritic = actor_critic_class( - self.env.num_obs, num_critic_obs, self.env.num_actions, **self.policy_cfg - ).to(self.device) - - alg_class = eval(self.cfg["algorithm_class_name"]) # PPO - - self.alg: PPO = alg_class(actor_critic, device=self.device, **self.alg_cfg) - self.num_steps_per_env = self.cfg["num_steps_per_env"] - self.save_interval = self.cfg["save_interval"] - - # init storage and model - self.alg.init_storage( - self.env.num_envs, - self.num_steps_per_env, - [self.env.num_obs], - [self.env.num_privileged_obs], - [self.env.num_actions], - ) - - # Log - self.log_dir = log_dir - self.writer = None - self.tot_timesteps = 0 - self.tot_time = 0 - self.current_learning_iteration = 0 - - _, _ = self.env.reset() - - def learn(self, num_learning_iterations, init_at_random_ep_len=False): - # initialize writer - if self.log_dir is not None and self.writer is None: - self.writer = SummaryWriter(log_dir=self.log_dir, flush_secs=10) - if init_at_random_ep_len: - self.env.episode_length_buf = torch.randint_like( - self.env.episode_length_buf, high=int(self.env.max_episode_length) - ) - obs = self.env.get_observations() - privileged_obs = self.env.get_privileged_observations() - critic_obs = privileged_obs if privileged_obs is not None else obs - obs, critic_obs = obs.to(self.device), critic_obs.to(self.device) - self.alg.actor_critic.train() # switch to train mode (for dropout for example) - - ep_infos = [] - rewbuffer = deque(maxlen=100) - costbuffer = deque(maxlen=100) - lenbuffer = deque(maxlen=100) - cur_reward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - cur_cost_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - cur_episode_length = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - - tot_iter = self.current_learning_iteration + num_learning_iterations - for it in range(self.current_learning_iteration, tot_iter): - start = time.time() - # Rollout - with torch.inference_mode(): - for i in range(self.num_steps_per_env): - # import ipdb; ipdb.set_trace() - actions = self.alg.act(obs, critic_obs) - obs, privileged_obs, rewards, dones, infos = self.env.step(actions.detach()) - - critic_obs = privileged_obs if privileged_obs is not None else obs - obs, critic_obs, rewards, dones = ( - obs.to(self.device), - critic_obs.to(self.device), - rewards.to(self.device), - dones.to(self.device), - ) - - self.alg.process_env_step(rewards, dones, infos) - - if self.log_dir is not None: - # Book keeping - if "log" in infos: - ep_infos.append(infos["log"]) - cur_reward_sum += rewards - cur_episode_length += 1 - new_ids = (dones > 0).nonzero(as_tuple=False) - rewbuffer.extend(cur_reward_sum[new_ids][:, 0].cpu().numpy().tolist()) - costbuffer.extend(cur_cost_sum[new_ids][:, 0].cpu().numpy().tolist()) - lenbuffer.extend(cur_episode_length[new_ids][:, 0].cpu().numpy().tolist()) - cur_reward_sum[new_ids] = 0 - cur_cost_sum[new_ids] = 0 - cur_episode_length[new_ids] = 0 - - stop = time.time() - collection_time = stop - start - - # Learning step - start = stop - self.alg.compute_returns(critic_obs) - - teleop_body_pos_upperbody_sigma = self.env.cfg.rewards.body_pos_upper_body_sigma - penalty_scale = self.env.unwrapped.penalty_scale - average_episode_length_for_reward_curriculum = self.env.unwrapped.average_episode_length - - mean_value_loss, mean_surrogate_loss = self.alg.update(epoch=it) - stop = time.time() - learn_time = stop - start - if self.log_dir is not None: - self.log(locals()) - - if it % self.save_interval == 0: - self.save(os.path.join(self.log_dir, "model_{}.pt".format(it))) - ep_infos.clear() - - self.current_learning_iteration += num_learning_iterations - self.save(os.path.join(self.log_dir, "model_{}.pt".format(self.current_learning_iteration))) - - def log(self, locs, width=80, pad=35): - - self.tot_timesteps += self.num_steps_per_env * self.env.num_envs - self.tot_time += locs["collection_time"] + locs["learn_time"] - iteration_time = locs["collection_time"] + locs["learn_time"] - - ep_string = f"" - if locs["ep_infos"]: - for key in locs["ep_infos"][0]: - infotensor = torch.tensor([], device=self.device) - for ep_info in locs["ep_infos"]: - # handle scalar and zero dimensional tensor infos - if not isinstance(ep_info[key], torch.Tensor): - ep_info[key] = torch.Tensor([ep_info[key]]) - if len(ep_info[key].shape) == 0: - ep_info[key] = ep_info[key].unsqueeze(0) - infotensor = torch.cat((infotensor, ep_info[key].to(self.device))) - value = torch.mean(infotensor) - self.writer.add_scalar("Episode/" + key, value, locs["it"]) - ep_string += f"""{f'Mean episode {key}:':>{pad}} {value:.4f}\n""" - mean_std = self.alg.actor_critic.std.mean() - fps = int(self.num_steps_per_env * self.env.num_envs / (locs["collection_time"] + locs["learn_time"])) - - self.writer.add_scalar( - "Episode/Average_episode_length_for_reward_curriculum", - locs["average_episode_length_for_reward_curriculum"], - locs["it"], - ) - self.writer.add_scalar("Episode/Penalty_scale", locs["penalty_scale"], locs["it"]) - self.writer.add_scalar( - "Episode/Teleop_body_pos_upperbody_sigma", locs["teleop_body_pos_upperbody_sigma"], locs["it"] - ) - - self.writer.add_scalar("Loss/value_function", locs["mean_value_loss"], locs["it"]) - self.writer.add_scalar("Loss/surrogate", locs["mean_surrogate_loss"], locs["it"]) - self.writer.add_scalar("Loss/learning_rate", self.alg.learning_rate, locs["it"]) - self.writer.add_scalar("Policy/mean_noise_std", mean_std.item(), locs["it"]) - self.writer.add_scalar("Perf/total_fps", fps, locs["it"]) - self.writer.add_scalar("Perf/collection time", locs["collection_time"], locs["it"]) - self.writer.add_scalar("Perf/learning_time", locs["learn_time"], locs["it"]) - - if len(locs["rewbuffer"]) > 0: - self.writer.add_scalar("Train/mean_reward", statistics.mean(locs["rewbuffer"]), locs["it"]) - self.writer.add_scalar("Train/mean_cost", statistics.mean(locs["costbuffer"]), locs["it"]) - self.writer.add_scalar("Train/mean_episode_length", statistics.mean(locs["lenbuffer"]), locs["it"]) - self.writer.add_scalar("Train/mean_reward/time", statistics.mean(locs["rewbuffer"]), self.tot_time) - self.writer.add_scalar("Train/mean_cost/time", statistics.mean(locs["costbuffer"]), self.tot_time) - self.writer.add_scalar("Train/mean_episode_length/time", statistics.mean(locs["lenbuffer"]), self.tot_time) - - str = f" \033[1m Learning iteration {locs['it']}/{self.current_learning_iteration + locs['num_learning_iterations']} \033[0m " - - if len(locs["rewbuffer"]) > 0: - log_string = ( - f"""{'#' * width}\n""" - f"""{str.center(width, ' ')}\n\n""" - f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ - 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" - f"""{'Value function loss:':>{pad}} {locs['mean_value_loss']:.4f}\n""" - f"""{'Surrogate loss:':>{pad}} {locs['mean_surrogate_loss']:.4f}\n""" - f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" - f"""{'Mean reward:':>{pad}} {statistics.mean(locs['rewbuffer']):.2f}\n""" - f"""{'Mean cost:':>{pad}} {statistics.mean(locs['costbuffer']):.2f}\n""" - f"""{'Mean episode length:':>{pad}} {statistics.mean(locs['lenbuffer']):.2f}\n""" - f"""{'Average_episode_length_for_reward_curriculum:':>{pad}} {locs['average_episode_length_for_reward_curriculum']:.6f}\n""" - f"""{'Penalty_scale:':>{pad}} {locs['penalty_scale']:.6f}\n""" - f"""{'Teleop_body_pos_upperbody_sigma:':>{pad}} {locs['teleop_body_pos_upperbody_sigma']:.6f}\n""" - ) - else: - log_string = ( - f"""{'#' * width}\n""" - f"""{str.center(width, ' ')}\n\n""" - f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ - 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" - f"""{'Value function loss:':>{pad}} {locs['mean_value_loss']:.4f}\n""" - f"""{'Surrogate loss:':>{pad}} {locs['mean_surrogate_loss']:.4f}\n""" - f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" - ) - - log_string += ep_string - - log_string += ( - f"""{'-' * width}\n""" - f"""{'Total timesteps:':>{pad}} {self.tot_timesteps}\n""" - f"""{'Iteration time:':>{pad}} {iteration_time:.2f}s\n""" - f"""{'Total time:':>{pad}} {self.tot_time:.2f}s\n""" - f"""{'ETA:':>{pad}} {self.format_seconds_to_human_readable(self.tot_time / (locs['it'] + 1) * ( - locs['num_learning_iterations'] - locs['it']))}\n""" - ) - - log_string += f"""path: {self.log_dir}\n""" - print("\r " + log_string, end="") - - def format_seconds_to_human_readable(self, total_seconds): - """Formats seconds into a human readable string with hours, minutes and seconds. - - Args: - total_seconds (float): Number of seconds to format - - Returns: - str: Formatted string in the format "Xh, Ym, Zs" where X=hours, Y=minutes, Z=seconds - """ - hours: int = total_seconds // 3600 - minutes: int = (total_seconds % 3600) // 60 - seconds: float = total_seconds % 60 - return f"{hours:.0f}h, {minutes:.0f}m, {seconds:.1f}s" - - def save(self, path, infos=None): - torch.save( - { - "model_state_dict": self.alg.actor_critic.state_dict(), - "optimizer_state_dict": self.alg.optimizer.state_dict(), - "iter": self.current_learning_iteration, - "infos": infos, - }, - path, - ) - - def load(self, path, load_optimizer=True): - loaded_dict = torch.load(path, map_location=self.device) - self.alg.actor_critic.load_state_dict(loaded_dict["model_state_dict"]) - if load_optimizer: - self.alg.optimizer.load_state_dict(loaded_dict["optimizer_state_dict"]) - self.current_learning_iteration = loaded_dict["iter"] - return loaded_dict["infos"] - - def get_inference_policy(self, device=None): - self.alg.actor_critic.eval() # switch to evaluation mode (dropout for example) - if device is not None: - self.alg.actor_critic.to(device) - return self.alg.actor_critic.act_inference diff --git a/third_party/rsl_rl/rsl_rl/storage/__init__.py b/third_party/rsl_rl/rsl_rl/storage/__init__.py deleted file mode 100755 index cbe7d80..0000000 --- a/third_party/rsl_rl/rsl_rl/storage/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -# Copyright 2021 ETH Zurich, NVIDIA CORPORATION -# SPDX-License-Identifier: BSD-3-Clause - -from .rollout_storage import RolloutStorage diff --git a/third_party/rsl_rl/rsl_rl/storage/rollout_storage.py b/third_party/rsl_rl/rsl_rl/storage/rollout_storage.py deleted file mode 100755 index fe1ce4f..0000000 --- a/third_party/rsl_rl/rsl_rl/storage/rollout_storage.py +++ /dev/null @@ -1,305 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: BSD-3-Clause -# -# Redistribution and use in source and binary forms, with or without -# modification, are permitted provided that the following conditions are met: -# -# 1. Redistributions of source code must retain the above copyright notice, this -# list of conditions and the following disclaimer. -# -# 2. Redistributions in binary form must reproduce the above copyright notice, -# this list of conditions and the following disclaimer in the documentation -# and/or other materials provided with the distribution. -# -# 3. Neither the name of the copyright holder nor the names of its -# contributors may be used to endorse or promote products derived from -# this software without specific prior written permission. -# -# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE -# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR -# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER -# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, -# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -# -# Copyright (c) 2021 ETH Zurich, Nikita Rudin - -import torch -import numpy as np - -from rsl_rl.utils import split_and_pad_trajectories - - -class RolloutStorage: - class Transition: - def __init__(self): - self.observations = None - self.critic_observations = None - self.actions = None - self.rewards = None - self.dones = None - self.values = None - self.actions_log_prob = None - self.action_mean = None - self.action_sigma = None - self.hidden_states = None - self.kin_dict = None - - def clear(self): - self.__init__() - - def __init__( - self, - num_envs, - num_transitions_per_env, - obs_shape, - privileged_obs_shape, - actions_shape, - kin_dict_info=None, - device="cpu", - ): - - self.device = device - - self.obs_shape = obs_shape - self.privileged_obs_shape = privileged_obs_shape - self.actions_shape = actions_shape - - # Core - self.observations = torch.zeros(num_transitions_per_env, num_envs, *obs_shape, device=self.device) - if privileged_obs_shape[0] is not None: - self.privileged_observations = torch.zeros( - num_transitions_per_env, num_envs, *privileged_obs_shape, device=self.device - ) - else: - self.privileged_observations = None - self.rewards = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) - self.actions = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) - self.dones = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device).byte() - if not kin_dict_info is None: - kin_dict_size = np.sum([v[-1][-1] for k, v in kin_dict_info.items()]) - self.kin_dict = torch.zeros(num_transitions_per_env, num_envs, kin_dict_size, device=self.device) - else: - self.kin_dict = None - - # For PPO - self.actions_log_prob = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) - self.values = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) - self.returns = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) - self.advantages = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) - self.mu = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) - self.sigma = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) - - self.num_transitions_per_env = num_transitions_per_env - self.num_envs = num_envs - - # rnn - self.saved_hidden_states_a = None - self.saved_hidden_states_c = None - - self.step = 0 - - def add_transitions(self, transition: Transition): - if self.step >= self.num_transitions_per_env: - raise AssertionError("Rollout buffer overflow") - self.observations[self.step].copy_(transition.observations) - if self.privileged_observations is not None: - self.privileged_observations[self.step].copy_(transition.critic_observations) - self.actions[self.step].copy_(transition.actions) - self.rewards[self.step].copy_(transition.rewards.view(-1, 1)) - self.dones[self.step].copy_(transition.dones.view(-1, 1)) - self.values[self.step].copy_(transition.values) - self.actions_log_prob[self.step].copy_(transition.actions_log_prob.view(-1, 1)) - self.mu[self.step].copy_(transition.action_mean) - self.sigma[self.step].copy_(transition.action_sigma) - if not self.kin_dict is None: - self.kin_dict[self.step].copy_(transition.kin_dict) - self._save_hidden_states(transition.hidden_states) - self.step += 1 - - def _save_hidden_states(self, hidden_states): - if hidden_states is None or hidden_states == (None, None): - return - # make a tuple out of GRU hidden state sto match the LSTM format - hid_a = hidden_states[0] if isinstance(hidden_states[0], tuple) else (hidden_states[0],) - hid_c = hidden_states[1] if isinstance(hidden_states[1], tuple) else (hidden_states[1],) - - # initialize if needed - if self.saved_hidden_states_a is None: - self.saved_hidden_states_a = [ - torch.zeros(self.observations.shape[0], *hid_a[i].shape, device=self.device) for i in range(len(hid_a)) - ] - self.saved_hidden_states_c = [ - torch.zeros(self.observations.shape[0], *hid_c[i].shape, device=self.device) for i in range(len(hid_c)) - ] - # copy the states - for i in range(len(hid_a)): - self.saved_hidden_states_a[i][self.step].copy_(hid_a[i]) - self.saved_hidden_states_c[i][self.step].copy_(hid_c[i]) - - def clear(self): - self.step = 0 - - def compute_returns(self, last_values, gamma, lam): - advantage = 0 - for step in reversed(range(self.num_transitions_per_env)): - if step == self.num_transitions_per_env - 1: - next_values = last_values - else: - next_values = self.values[step + 1] - next_is_not_terminal = 1.0 - self.dones[step].float() - delta = self.rewards[step] + next_is_not_terminal * gamma * next_values - self.values[step] - advantage = delta + next_is_not_terminal * gamma * lam * advantage - self.returns[step] = advantage + self.values[step] - - # Compute and normalize the advantages - self.advantages = self.returns - self.values - self.advantages = (self.advantages - self.advantages.mean()) / (self.advantages.std() + 1e-8) - - def get_statistics(self): - done = self.dones - done[-1] = 1 - flat_dones = done.permute(1, 0, 2).reshape(-1, 1) - done_indices = torch.cat( - (flat_dones.new_tensor([-1], dtype=torch.int64), flat_dones.nonzero(as_tuple=False)[:, 0]) - ) - trajectory_lengths = done_indices[1:] - done_indices[:-1] - return trajectory_lengths.float().mean(), self.rewards.mean() - - def mini_batch_generator(self, num_mini_batches, num_epochs=8): - batch_size = self.num_envs * self.num_transitions_per_env - mini_batch_size = batch_size // num_mini_batches - indices = torch.randperm(num_mini_batches * mini_batch_size, requires_grad=False, device=self.device) - - observations = self.observations.flatten(0, 1) - if self.privileged_observations is not None: - critic_observations = self.privileged_observations.flatten(0, 1) - else: - critic_observations = observations - - actions = self.actions.flatten(0, 1) - values = self.values.flatten(0, 1) - returns = self.returns.flatten(0, 1) - old_actions_log_prob = self.actions_log_prob.flatten(0, 1) - advantages = self.advantages.flatten(0, 1) - old_mu = self.mu.flatten(0, 1) - old_sigma = self.sigma.flatten(0, 1) - if not self.kin_dict is None: - kin_dict = self.kin_dict.flatten(0, 1) - # import ipdb; ipdb.set_trace() - for epoch in range(num_epochs): - for i in range(num_mini_batches): - - start = i * mini_batch_size - end = (i + 1) * mini_batch_size - stop = (i + 1) * mini_batch_size - batch_idx = indices[start:end] - - obs_batch = observations[batch_idx] - critic_observations_batch = critic_observations[batch_idx] - actions_batch = actions[batch_idx] - - target_values_batch = values[batch_idx] - returns_batch = returns[batch_idx] - old_actions_log_prob_batch = old_actions_log_prob[batch_idx] - advantages_batch = advantages[batch_idx] - old_mu_batch = old_mu[batch_idx] - old_sigma_batch = old_sigma[batch_idx] - - obs_squential = observations[start:end] - # import ipdb; ipdb.set_trace() - if not self.kin_dict is None: - kin_dict_batch = kin_dict[batch_idx] - yield obs_batch, critic_observations_batch, actions_batch, target_values_batch, advantages_batch, returns_batch, old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, ( - None, - None, - ), None, obs_squential, kin_dict_batch - - else: - - yield obs_batch, critic_observations_batch, actions_batch, target_values_batch, advantages_batch, returns_batch, old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, ( - None, - None, - ), None, obs_squential, None - - # for RNNs only - def reccurent_mini_batch_generator(self, num_mini_batches, num_epochs=8): - - padded_obs_trajectories, trajectory_masks = split_and_pad_trajectories(self.observations, self.dones) - if self.privileged_observations is not None: - padded_critic_obs_trajectories, _ = split_and_pad_trajectories(self.privileged_observations, self.dones) - else: - padded_critic_obs_trajectories = padded_obs_trajectories - - mini_batch_size = self.num_envs // num_mini_batches - if not self.kin_dict is None: - kin_dict = self.kin_dict # flatten(0, 1) - - for ep in range(num_epochs): - first_traj = 0 - for i in range(num_mini_batches): - start = i * mini_batch_size - stop = (i + 1) * mini_batch_size - - dones = self.dones.squeeze(-1) - last_was_done = torch.zeros_like(dones, dtype=torch.bool) - last_was_done[1:] = dones[:-1] - last_was_done[0] = True - trajectories_batch_size = torch.sum(last_was_done[:, start:stop]) - last_traj = first_traj + trajectories_batch_size - - # import ipdb; ipdb.set_trace() - - masks_batch = trajectory_masks[:, first_traj:last_traj] - obs_batch = padded_obs_trajectories[:, first_traj:last_traj] - critic_obs_batch = padded_critic_obs_trajectories[:, first_traj:last_traj] - - actions_batch = self.actions[:, start:stop] - old_mu_batch = self.mu[:, start:stop] - old_sigma_batch = self.sigma[:, start:stop] - returns_batch = self.returns[:, start:stop] - advantages_batch = self.advantages[:, start:stop] - values_batch = self.values[:, start:stop] - old_actions_log_prob_batch = self.actions_log_prob[:, start:stop] - - obs_squential = self.observations[:, start:stop] - # reshape to [num_envs, time, num layers, hidden dim] (original shape: [time, num_layers, num_envs, hidden_dim]) - # then take only time steps after dones (flattens num envs and time dimensions), - # take a batch of trajectories and finally reshape back to [num_layers, batch, hidden_dim] - last_was_done = last_was_done.permute(1, 0) - hid_a_batch = [ - saved_hidden_states.permute(2, 0, 1, 3)[last_was_done][first_traj:last_traj] - .transpose(1, 0) - .contiguous() - for saved_hidden_states in self.saved_hidden_states_a - ] - hid_c_batch = [ - saved_hidden_states.permute(2, 0, 1, 3)[last_was_done][first_traj:last_traj] - .transpose(1, 0) - .contiguous() - for saved_hidden_states in self.saved_hidden_states_c - ] - # remove the tuple for GRU - hid_a_batch = hid_a_batch[0] if len(hid_a_batch) == 1 else hid_a_batch - hid_c_batch = hid_c_batch[0] if len(hid_c_batch) == 1 else hid_c_batch - - if not self.kin_dict is None: - kin_dict_batch = kin_dict[:, start:stop] - # print("obs_batch", obs_batch.shape) - yield obs_batch, critic_obs_batch, actions_batch, values_batch, advantages_batch, returns_batch, old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, ( - hid_a_batch, - hid_c_batch, - ), masks_batch, obs_squential, kin_dict_batch - - else: - print("obs_batch", obs_batch.shape) - yield obs_batch, critic_obs_batch, actions_batch, values_batch, advantages_batch, returns_batch, old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, ( - hid_a_batch, - hid_c_batch, - ), masks_batch, obs_squential, None - - first_traj = last_traj diff --git a/third_party/rsl_rl/rsl_rl/utils/__init__.py b/third_party/rsl_rl/rsl_rl/utils/__init__.py deleted file mode 100755 index 1b505f3..0000000 --- a/third_party/rsl_rl/rsl_rl/utils/__init__.py +++ /dev/null @@ -1,31 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: BSD-3-Clause -# -# Redistribution and use in source and binary forms, with or without -# modification, are permitted provided that the following conditions are met: -# -# 1. Redistributions of source code must retain the above copyright notice, this -# list of conditions and the following disclaimer. -# -# 2. Redistributions in binary form must reproduce the above copyright notice, -# this list of conditions and the following disclaimer in the documentation -# and/or other materials provided with the distribution. -# -# 3. Neither the name of the copyright holder nor the names of its -# contributors may be used to endorse or promote products derived from -# this software without specific prior written permission. -# -# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE -# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR -# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER -# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, -# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -# -# Copyright (c) 2021 ETH Zurich, Nikita Rudin - -from .utils import split_and_pad_trajectories, unpad_trajectories \ No newline at end of file diff --git a/third_party/rsl_rl/rsl_rl/utils/utils.py b/third_party/rsl_rl/rsl_rl/utils/utils.py deleted file mode 100755 index b6affab..0000000 --- a/third_party/rsl_rl/rsl_rl/utils/utils.py +++ /dev/null @@ -1,71 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: BSD-3-Clause -# -# Redistribution and use in source and binary forms, with or without -# modification, are permitted provided that the following conditions are met: -# -# 1. Redistributions of source code must retain the above copyright notice, this -# list of conditions and the following disclaimer. -# -# 2. Redistributions in binary form must reproduce the above copyright notice, -# this list of conditions and the following disclaimer in the documentation -# and/or other materials provided with the distribution. -# -# 3. Neither the name of the copyright holder nor the names of its -# contributors may be used to endorse or promote products derived from -# this software without specific prior written permission. -# -# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE -# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR -# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER -# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, -# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -# -# Copyright (c) 2021 ETH Zurich, Nikita Rudin - -import torch - -def split_and_pad_trajectories(tensor, dones): - """ Splits trajectories at done indices. Then concatenates them and padds with zeros up to the length og the longest trajectory. - Returns masks corresponding to valid parts of the trajectories - Example: - Input: [ [a1, a2, a3, a4 | a5, a6], - [b1, b2 | b3, b4, b5 | b6] - ] - - Output:[ [a1, a2, a3, a4], | [ [True, True, True, True], - [a5, a6, 0, 0], | [True, True, False, False], - [b1, b2, 0, 0], | [True, True, False, False], - [b3, b4, b5, 0], | [True, True, True, False], - [b6, 0, 0, 0] | [True, False, False, False], - ] | ] - - Assumes that the inputy has the following dimension order: [time, number of envs, aditional dimensions] - """ - dones = dones.clone() - dones[-1] = 1 - # Permute the buffers to have order (num_envs, num_transitions_per_env, ...), for correct reshaping - flat_dones = dones.transpose(1, 0).reshape(-1, 1) - - # Get length of trajectory by counting the number of successive not done elements - done_indices = torch.cat((flat_dones.new_tensor([-1], dtype=torch.int64), flat_dones.nonzero()[:, 0])) - trajectory_lengths = done_indices[1:] - done_indices[:-1] - trajectory_lengths_list = trajectory_lengths.tolist() - # Extract the individual trajectories - trajectories = torch.split(tensor.transpose(1, 0).flatten(0, 1),trajectory_lengths_list) - padded_trajectories = torch.nn.utils.rnn.pad_sequence(trajectories) - - - trajectory_masks = trajectory_lengths > torch.arange(0, tensor.shape[0], device=tensor.device).unsqueeze(1) - return padded_trajectories, trajectory_masks - -def unpad_trajectories(trajectories, masks): - """ Does the inverse operation of split_and_pad_trajectories() - """ - # Need to transpose before and after the masking to have proper reshaping - return trajectories.transpose(1, 0)[masks.transpose(1, 0)].view(-1, trajectories.shape[0], trajectories.shape[-1]).transpose(1, 0) \ No newline at end of file diff --git a/third_party/rsl_rl/setup.py b/third_party/rsl_rl/setup.py deleted file mode 100755 index 7851872..0000000 --- a/third_party/rsl_rl/setup.py +++ /dev/null @@ -1,13 +0,0 @@ -from setuptools import setup, find_packages - -setup( - name="rsl_rl", - version="1.0.2", - author="Nikita Rudin", - author_email="rudinn@ethz.ch", - license="BSD-3-Clause", - packages=find_packages(), - description="Fast and simple RL algorithms implemented in pytorch", - python_requires=">=3.6", - install_requires=["torch>=1.4.0", "torchvision>=0.5.0", "numpy>=1.16.4"], -)