fix: [Breaking change] code structure change

* Separate and decouple different environments
* Simplified the process of adding obs to custom environments
* Fixed some bugs using local rsl_rl
* Place the script in the subdivided folder

.gitignore

@@ -33,4 +33,4 @@ docker/artifacts/
 _isaac_sim*
 _repo
 _build
-.lastformat
+.lastformat

.pre-commit-config.yaml

@@ -47,4 +47,4 @@ repos:
       - id: rst-backticks
       - id: rst-directive-colons
       - id: rst-inline-touching-normal
-exclude: ^(exts/robot_lab/amp_utils/|.vscode/)
+exclude: ^(exts/robot_lab/robot_lab/third_party/|.vscode/)

README.md

@@ -14,13 +14,6 @@
 
 If you want to run policy in gazebo or real robot, please use [rl_sar](https://github.com/fan-ziqi/rl_sar) project.
 
-Todo:
-
-- [x] AMP training
-- [ ] VAE training code
-- [x] Sim to Sim transfer(Gazebo)
-- [ ] Sim to Real transfer(Unitree A1)
-
 [Click to discuss on Discord](https://discord.gg/vmVjkhVugU)
 
 ## Installation
@@ -84,45 +77,45 @@ FFTAI GR1T1
 
 ```bash
 # Train
-python scripts/rsl_rl/train.py --task RobotLab-Isaac-Velocity-Flat-FFTAI-GR1T1-v0 --headless
+python scripts/rsl_rl/base/train.py --task RobotLab-Isaac-Velocity-Flat-FFTAI-GR1T1-v0 --headless
 # Play
-python scripts/rsl_rl/play.py --task RobotLab-Isaac-Velocity-Flat-FFTAI-GR1T1-v0
+python scripts/rsl_rl/base/play.py --task RobotLab-Isaac-Velocity-Flat-FFTAI-GR1T1-v0
 ```
 
 Anymal D
 
 ```bash
 # Train
-python scripts/rsl_rl/train.py --task RobotLab-Isaac-Velocity-Flat-Anymal-D-v0 --headless
+python scripts/rsl_rl/base/train.py --task RobotLab-Isaac-Velocity-Flat-Anymal-D-v0 --headless
 # Play
-python scripts/rsl_rl/play.py --task RobotLab-Isaac-Velocity-Flat-Anymal-D-v0
+python scripts/rsl_rl/base/play.py --task RobotLab-Isaac-Velocity-Flat-Anymal-D-v0
 ```
 
 Unitree A1
 
 ```bash
 # Train
-python scripts/rsl_rl/train.py --task RobotLab-Isaac-Velocity-Flat-Unitree-A1-v0 --headless
+python scripts/rsl_rl/base/train.py --task RobotLab-Isaac-Velocity-Flat-Unitree-A1-v0 --headless
 # Play
-python scripts/rsl_rl/play.py --task RobotLab-Isaac-Velocity-Flat-Unitree-A1-v0
+python scripts/rsl_rl/base/play.py --task RobotLab-Isaac-Velocity-Flat-Unitree-A1-v0
 ```
 
 Unitree Go2W (Unvalible for now)
 
 ```bash
 # Train
-python scripts/rsl_rl/train.py --task RobotLab-Isaac-Velocity-Flat-Unitree-Go2W-v0 --headless
+python scripts/rsl_rl/base/train.py --task RobotLab-Isaac-Velocity-Flat-Unitree-Go2W-v0 --headless
 # Play
-python scripts/rsl_rl/play.py --task RobotLab-Isaac-Velocity-Flat-Unitree-Go2W-v0
+python scripts/rsl_rl/base/play.py --task RobotLab-Isaac-Velocity-Flat-Unitree-Go2W-v0
 ```
 
 Unitree H1
 
 ```bash
 # Train
-python scripts/rsl_rl/train.py --task RobotLab-Isaac-Velocity-Flat-Unitree-H1-v0 --headless
+python scripts/rsl_rl/base/train.py --task RobotLab-Isaac-Velocity-Flat-Unitree-H1-v0 --headless
 # Play
-python scripts/rsl_rl/play.py --task RobotLab-Isaac-Velocity-Flat-Unitree-H1-v0
+python scripts/rsl_rl/base/play.py --task RobotLab-Isaac-Velocity-Flat-Unitree-H1-v0
 ```
 
 The above configs are flat, you can change Flat to Rough
@@ -142,13 +135,13 @@ Unitree A1
 
 ```bash
 # Retarget motion files
-python exts/robot_lab/amp_utils/scripts/retarget_kp_motions.py
+python exts/robot_lab/robot_lab/third_party/amp_utils/scripts/retarget_kp_motions.py
 # Replay AMP data
-python scripts/rsl_rl/replay_amp_data.py --task RobotLab-Isaac-Velocity-Flat-Amp-Unitree-A1-v0
+python scripts/rsl_rl/amp/replay_amp_data.py --task RobotLab-Isaac-Velocity-Flat-Amp-Unitree-A1-v0
 # Train
-python scripts/rsl_rl/train_amp.py --task RobotLab-Isaac-Velocity-Flat-Amp-Unitree-A1-v0 --headless
+python scripts/rsl_rl/amp/train.py --task RobotLab-Isaac-Velocity-Flat-Amp-Unitree-A1-v0 --headless
 # Play
-python scripts/rsl_rl/play_amp.py --task RobotLab-Isaac-Velocity-Flat-Amp-Unitree-A1-v0
+python scripts/rsl_rl/amp/play.py --task RobotLab-Isaac-Velocity-Flat-Amp-Unitree-A1-v0
 ```
 
 ## Add your own robot

exts/robot_lab/config/extension.toml

@@ -1,16 +1,16 @@
 [package]
 
 # Semantic Versioning is used: https://semver.org/
-version = "0.4.0"
+version = "1.0"
 
 # Description
 category = "isaaclab"
 readme  = "README.md"
 
-title = "Real Robot Extension"
+title = "A generic robot RL library based on IsaacLab"
 author = "Ziqi Fan"
 maintainer = "Ziqi Fan"
-description="Real Robot Training for Isaac Lab"
+description="A generic robot RL library based on IsaacLab"
 repository = "https://github.com/fan-ziqi/robot_lab.git"
 keywords = ["extension", "robot_lab", "isaaclab"]
 

exts/robot_lab/robot_lab/assets/fftai.py

@@ -7,12 +7,12 @@
 Reference: https://github.com/FFTAI
 """
 
-from robot_lab.assets import ISAACLAB_ASSETS_DATA_DIR
-
 import omni.isaac.lab.sim as sim_utils
 from omni.isaac.lab.actuators import ImplicitActuatorCfg
 from omni.isaac.lab.assets.articulation import ArticulationCfg
 
+from robot_lab.assets import ISAACLAB_ASSETS_DATA_DIR
+
 ##
 # Configuration
 ##

exts/robot_lab/robot_lab/assets/unitree.py

@@ -8,12 +8,12 @@
 Reference: https://github.com/unitreerobotics/unitree_ros
 """
 
-from robot_lab.assets import ISAACLAB_ASSETS_DATA_DIR
-
 import omni.isaac.lab.sim as sim_utils
 from omni.isaac.lab.actuators import DCMotorCfg
 from omni.isaac.lab.assets.articulation import ArticulationCfg
 
+from robot_lab.assets import ISAACLAB_ASSETS_DATA_DIR
+
 ##
 # Configuration
 ##

exts/robot_lab/robot_lab/tasks/locomotion/velocity/config/anymal_d/rough_env_cfg.py

@@ -24,14 +24,12 @@ def __post_init__(self):
         # ------------------------------Observations------------------------------
         self.observations.policy.base_lin_vel = None
         self.observations.policy.height_scan = None
-        self.observations.AMP = None
 
         # ------------------------------Actions------------------------------
         # reduce action scale
         self.actions.joint_pos.scale = 0.5
 
         # ------------------------------Events------------------------------
-        self.events.reset_amp = None
         self.events.add_base_mass.params["mass_distribution_params"] = (-5.0, 5.0)
         self.events.add_base_mass.params["asset_cfg"].body_names = ["base"]
         self.events.base_external_force_torque.params["asset_cfg"].body_names = ["base"]

exts/robot_lab/robot_lab/tasks/locomotion/velocity/config/fftai_gr1t1/rough_env_cfg.py

@@ -1,6 +1,6 @@
-import robot_lab.tasks.locomotion.velocity.mdp as mdp
 from omni.isaac.lab.utils import configclass
 
+import robot_lab.tasks.locomotion.velocity.mdp as mdp
 from robot_lab.tasks.locomotion.velocity.velocity_env_cfg import LocomotionVelocityRoughEnvCfg
 
 ##
@@ -26,14 +26,12 @@ def __post_init__(self):
         # ------------------------------Observations------------------------------
         # self.observations.policy.base_lin_vel = None
         # self.observations.policy.height_scan = None
-        self.observations.AMP = None
 
         # ------------------------------Actions------------------------------
         # reduce action scale
         # self.actions.joint_pos.scale = 1.0
 
         # ------------------------------Events------------------------------
-        self.events.reset_amp = None
         self.events.base_external_force_torque.params["asset_cfg"].body_names = ["base"]
         self.events.reset_robot_joints.params["position_range"] = (1.0, 1.0)
         self.events.reset_base.params = {

exts/robot_lab/robot_lab/tasks/locomotion/velocity/config/unitree_a1/rough_env_cfg.py

@@ -35,15 +35,13 @@ def __post_init__(self):
         self.observations.policy.joint_vel.scale = 0.05
         self.observations.policy.base_lin_vel = None
         self.observations.policy.height_scan = None
-        self.observations.AMP = None
 
         # ------------------------------Actions------------------------------
         # reduce action scale
         self.actions.joint_pos.scale = 0.25
         self.actions.joint_pos.clip = {".*": (-100, 100)}
 
         # ------------------------------Events------------------------------
-        self.events.reset_amp = None
         self.events.add_base_mass.params["mass_distribution_params"] = (-1.0, 3.0)
         self.events.add_base_mass.params["asset_cfg"].body_names = ["trunk"]
         self.events.base_external_force_torque.params["asset_cfg"].body_names = ["trunk"]

exts/robot_lab/robot_lab/tasks/locomotion/velocity/config/unitree_a1_amp/__init__.py

@@ -8,7 +8,7 @@
 
 gym.register(
     id="RobotLab-Isaac-Velocity-Flat-Amp-Unitree-A1-v0",
-    entry_point="robot_lab.tasks.locomotion.velocity.manager_based_rl_amp_env:ManagerBasedRLAmpEnv",
+    entry_point="robot_lab.tasks.locomotion.velocity.config.unitree_a1_amp.env.manager_based_rl_amp_env:ManagerBasedRLAmpEnv",
     disable_env_checker=True,
     kwargs={
         "env_cfg_entry_point": flat_env_cfg.UnitreeA1AmpFlatEnvCfg,
@@ -18,7 +18,7 @@
 
 gym.register(
     id="RobotLab-Isaac-Velocity-Rough-Amp-Unitree-A1-v0",
-    entry_point="robot_lab.tasks.locomotion.velocity.manager_based_rl_amp_env:ManagerBasedRLAmpEnv",
+    entry_point="robot_lab.tasks.locomotion.velocity.config.unitree_a1_amp.env.manager_based_rl_amp_env:ManagerBasedRLAmpEnv",
     disable_env_checker=True,
     kwargs={
         "env_cfg_entry_point": rough_env_cfg.UnitreeA1AmpRoughEnvCfg,

exts/robot_lab/robot_lab/tasks/locomotion/velocity/mdp/events.py → exts/robot_lab/robot_lab/tasks/locomotion/velocity/config/unitree_a1_amp/env/events.py

@@ -2,12 +2,12 @@
 
 import torch
 
-from robot_lab.tasks.locomotion.velocity.manager_based_rl_amp_env import ManagerBasedRLAmpEnv
-
 from omni.isaac.lab.assets import Articulation, RigidObject
 from omni.isaac.lab.managers import SceneEntityCfg
 from omni.isaac.lab.utils.math import quat_rotate
 
+from robot_lab.tasks.locomotion.velocity.config.unitree_a1_amp.env.manager_based_rl_amp_env import ManagerBasedRLAmpEnv
+
 
 def reset_amp(
     env: ManagerBasedRLAmpEnv,

exts/robot_lab/robot_lab/tasks/locomotion/velocity/manager_based_rl_amp_env.py → exts/robot_lab/robot_lab/tasks/locomotion/velocity/config/unitree_a1_amp/env/manager_based_rl_amp_env.py

@@ -9,13 +9,13 @@
 import gymnasium as gym
 import torch
 
-from amp_utils.kinematics import urdf
-from rsl_rl_extension.datasets.motion_loader import AMPLoader
-
 from omni.isaac.lab.envs.common import VecEnvStepReturn
 from omni.isaac.lab.envs.manager_based_rl_env import ManagerBasedRLEnv
 from omni.isaac.lab.envs.manager_based_rl_env_cfg import ManagerBasedRLEnvCfg
 
+from robot_lab.third_party.amp_utils.kinematics import urdf
+from robot_lab.third_party.rsl_rl_amp.datasets.motion_loader import AMPLoader
+
 
 class ManagerBasedRLAmpEnv(ManagerBasedRLEnv, gym.Env):
     def __init__(self, cfg: ManagerBasedRLEnvCfg, render_mode: str | None = None, **kwargs):

exts/robot_lab/robot_lab/tasks/locomotion/velocity/config/unitree_a1_amp/rough_env_cfg.py

@@ -1,10 +1,13 @@
 import glob
 
-from amp_utils import AMP_UTILS_DIR
-
+from omni.isaac.lab.managers import EventTermCfg as EventTerm
+from omni.isaac.lab.managers import ObservationTermCfg as ObsTerm
 from omni.isaac.lab.utils import configclass
 
-from robot_lab.tasks.locomotion.velocity.velocity_env_cfg import LocomotionVelocityRoughEnvCfg
+import robot_lab.tasks.locomotion.velocity.mdp as mdp
+from robot_lab.tasks.locomotion.velocity.config.unitree_a1_amp.env.events import reset_amp
+from robot_lab.tasks.locomotion.velocity.velocity_env_cfg import LocomotionVelocityRoughEnvCfg, create_obsgroup_class
+from robot_lab.third_party.amp_utils import AMP_UTILS_DIR
 
 ##
 # Pre-defined configs
@@ -40,6 +43,14 @@ def __post_init__(self):
         self.observations.policy.base_lin_vel = None
         self.observations.policy.base_ang_vel = None
         self.observations.policy.height_scan = None
+        self.observations.AMP = create_obsgroup_class('AMPCfg',{
+            'base_pos_z': ObsTerm(func=mdp.base_pos_z),
+            'base_lin_vel': ObsTerm(func=mdp.base_lin_vel),
+            'base_ang_vel': ObsTerm(func=mdp.base_ang_vel),
+            'joint_pos': ObsTerm(func=mdp.joint_pos),
+            'joint_vel': ObsTerm(func=mdp.joint_vel),
+        }, enable_corruption=True, concatenate_terms=True)()
+
 
         # ------------------------------Actions------------------------------
         # reduce action scale
@@ -55,6 +66,7 @@ def __post_init__(self):
         self.events.randomize_actuator_gains = None
         self.events.randomize_joint_parameters = None
         self.events.push_robot = None
+        self.events.reset_amp = EventTerm(func=reset_amp, mode="reset")
 
         # ------------------------------Rewards------------------------------
         # General

exts/robot_lab/robot_lab/tasks/locomotion/velocity/config/unitree_go2w/rough_env_cfg.py

@@ -33,15 +33,13 @@ def __post_init__(self):
         self.observations.policy.joint_vel.scale = 0.05
         self.observations.policy.base_lin_vel = None
         self.observations.policy.height_scan = None
-        self.observations.AMP = None
 
         # ------------------------------Actions------------------------------
         # reduce action scale
         self.actions.joint_pos.scale = 0.25
         self.actions.joint_pos.clip = {".*": (-100, 100)}
 
         # ------------------------------Events------------------------------
-        self.events.reset_amp = None
         self.events.add_base_mass.params["mass_distribution_params"] = (-1.0, 3.0)
         self.events.add_base_mass.params["asset_cfg"].body_names = ["base"]
         self.events.base_external_force_torque.params["asset_cfg"].body_names = ["base"]

exts/robot_lab/robot_lab/tasks/locomotion/velocity/config/unitree_h1/rough_env_cfg.py

@@ -1,6 +1,6 @@
-import robot_lab.tasks.locomotion.velocity.mdp as mdp
 from omni.isaac.lab.utils import configclass
 
+import robot_lab.tasks.locomotion.velocity.mdp as mdp
 from robot_lab.tasks.locomotion.velocity.velocity_env_cfg import LocomotionVelocityRoughEnvCfg
 
 ##
@@ -26,12 +26,10 @@ def __post_init__(self):
         # ------------------------------Observations------------------------------
         # self.observations.policy.base_lin_vel = None
         # self.observations.policy.height_scan = None
-        self.observations.AMP = None
 
         # ------------------------------Actions------------------------------
 
         # ------------------------------Events------------------------------
-        self.events.reset_amp = None
         self.events.base_external_force_torque.params["asset_cfg"].body_names = [".*torso_link"]
         self.events.reset_robot_joints.params["position_range"] = (1.0, 1.0)
         self.events.reset_base.params = {

exts/robot_lab/robot_lab/tasks/locomotion/velocity/mdp/__init__.py

@@ -1,8 +1,7 @@
 """This sub-module contains the functions that are specific to the locomotion environments."""
 
 from omni.isaac.lab.envs.mdp import *  # noqa: F401, F403
+from omni.isaac.lab_tasks.manager_based.locomotion.velocity.mdp import *  # noqa: F401, F403
 
 from .actions import *  # noqa: F401, F403
-from .curriculums import *  # noqa: F401, F403
-from .events import *  # noqa: F401, F403
 from .rewards import *  # noqa: F401, F403