Introducing locomotion module

source/uwlab_tasks/uwlab_tasks/manager_based/locomotion/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) 2024-2025, The UW Lab Project Developers.
+# All Rights Reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Locomotion environments for legged robots."""

source/uwlab_tasks/uwlab_tasks/manager_based/locomotion/advance_skills/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) 2024-2025, The UW Lab Project Developers.
+# All Rights Reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Advanced Skill Environments.
+Reference:
+    https://github.com/leggedrobotics/legged_gym
+"""

source/uwlab_tasks/uwlab_tasks/manager_based/locomotion/advance_skills/advance_skills_base_env.py

@@ -0,0 +1,315 @@
+# Copyright (c) 2024-2025, The UW Lab Project Developers.
+# All Rights Reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from dataclasses import MISSING
+
+import isaaclab.sim as sim_utils
+from isaaclab.assets import ArticulationCfg, AssetBaseCfg
+from isaaclab.envs import ViewerCfg
+from isaaclab.managers import CurriculumTermCfg as CurrTerm
+from isaaclab.managers import EventTermCfg as EventTerm
+from isaaclab.managers import ObservationGroupCfg as ObsGroup
+from isaaclab.managers import ObservationTermCfg as ObsTerm
+from isaaclab.managers import RewardTermCfg as RewTerm
+from isaaclab.managers import SceneEntityCfg
+from isaaclab.managers import TerminationTermCfg as DoneTerm
+from isaaclab.sensors import ContactSensorCfg, RayCasterCfg, patterns
+from isaaclab.utils import configclass
+from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR, ISAACLAB_NUCLEUS_DIR
+from isaaclab.utils.noise import AdditiveUniformNoiseCfg as Unoise
+from uwlab.envs.data_manager_based_rl import DataManagerBasedRLEnvCfg
+from uwlab.scene import InteractiveSceneCfg
+from uwlab.terrains import TerrainGeneratorCfg, TerrainImporterCfg
+
+import uwlab_tasks.manager_based.locomotion.advance_skills.mdp as mdp
+
+
+@configclass
+class SceneCfg(InteractiveSceneCfg):
+    """ "Configuration for the terrain scene with a legged robot."""
+
+    # ground terrain
+    terrain = TerrainImporterCfg(
+        prim_path="/World/ground",
+        terrain_type="generator",
+        terrain_generator=TerrainGeneratorCfg(
+            size=(10.0, 10.0),
+            border_width=20.0,
+            num_rows=10,
+            num_cols=20,
+            horizontal_scale=0.1,
+            vertical_scale=0.005,
+            slope_threshold=0.75,
+            use_cache=False,
+            sub_terrains={},
+        ),
+        max_init_terrain_level=5,
+        collision_group=-1,
+        physics_material=sim_utils.RigidBodyMaterialCfg(
+            friction_combine_mode="multiply",
+            restitution_combine_mode="multiply",
+            static_friction=1.0,
+            dynamic_friction=1.0,
+        ),
+        visual_material=sim_utils.MdlFileCfg(
+            mdl_path=f"{ISAACLAB_NUCLEUS_DIR}/Materials/TilesMarbleSpiderWhiteBrickBondHoned/TilesMarbleSpiderWhiteBrickBondHoned.mdl",
+            project_uvw=True,
+            texture_scale=(0.25, 0.25),
+        ),
+        debug_vis=False,
+    )
+
+    # lights
+    sky_light = AssetBaseCfg(
+        prim_path="/World/skyLight",
+        spawn=sim_utils.DomeLightCfg(
+            intensity=750.0,
+            texture_file=f"{ISAAC_NUCLEUS_DIR}/Materials/Textures/Skies/PolyHaven/kloofendal_43d_clear_puresky_4k.hdr",
+        ),
+    )
+
+    # robots
+    robot: ArticulationCfg = MISSING  # type: ignore
+
+    # sensors
+    height_scanner = RayCasterCfg(
+        prim_path="{ENV_REGEX_NS}/Robot/base",
+        offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 20.0)),
+        attach_yaw_only=True,
+        pattern_cfg=patterns.GridPatternCfg(resolution=0.1, size=(1.6, 1.0)),
+        debug_vis=False,
+        mesh_prim_paths=["/World/ground"],
+    )
+    contact_forces = ContactSensorCfg(
+        prim_path="{ENV_REGEX_NS}/Robot/.*", history_length=3, track_air_time=True, debug_vis=True
+    )
+
+
+@configclass
+class ActionsCfg:
+    """Actions for the MDP."""
+
+    joint_pos = mdp.JointPositionActionCfg(asset_name="robot", joint_names=[".*"], scale=0.5, use_default_offset=True)
+
+
+@configclass
+class CommandsCfg:
+    "Command specifications for the MDP."
+    goal_point = mdp.TerrainBasedPose2dCommandCfg(
+        asset_name="robot",
+        resampling_time_range=(10.0, 10.0),
+        simple_heading=False,
+        debug_vis=True,
+        ranges=mdp.TerrainBasedPose2dCommandCfg.Ranges(
+            heading=(-3.14, 3.14),
+        ),
+    )
+
+
+@configclass
+class ObservationsCfg:
+    """Observations for the MDP"""
+
+    @configclass
+    class PolicyCfg(ObsGroup):
+        base_lin_vel = ObsTerm(func=mdp.base_lin_vel, noise=Unoise(n_min=-0.1, n_max=0.1))
+        base_ang_vel = ObsTerm(func=mdp.base_ang_vel, noise=Unoise(n_min=-0.2, n_max=0.2))
+        proj_gravity = ObsTerm(func=mdp.projected_gravity, noise=Unoise(n_min=-0.05, n_max=0.05))
+        goal_point_commands = ObsTerm(func=mdp.generated_commands, params={"command_name": "goal_point"})
+        time_left = ObsTerm(func=mdp.time_left)
+        joint_pos = ObsTerm(func=mdp.joint_pos, noise=Unoise(n_min=-0.01, n_max=0.01))
+        joint_vel = ObsTerm(func=mdp.joint_vel, noise=Unoise(n_min=-1.5, n_max=1.5))
+        last_actions = ObsTerm(func=mdp.last_action)
+        height_scan = ObsTerm(
+            func=mdp.height_scan,
+            params={"sensor_cfg": SceneEntityCfg("height_scanner")},
+            noise=Unoise(n_min=-0.1, n_max=0.1),
+            clip=(-1.0, 1.0),
+        )
+
+        def __post_init__(self):
+            self.enable_corruption = True
+            self.concatenate_terms = True
+
+    policy: PolicyCfg = PolicyCfg()
+
+
+@configclass
+class EventsCfg:
+    # startup
+    physical_material = EventTerm(
+        func=mdp.randomize_rigid_body_material,  # type: ignore
+        mode="startup",
+        params={
+            "asset_cfg": SceneEntityCfg("robot", body_names=".*"),
+            "static_friction_range": (0.8, 0.8),
+            "dynamic_friction_range": (0.6, 0.6),
+            "restitution_range": (0.0, 0.0),
+            "num_buckets": 64,
+        },
+    )
+
+    add_base_mass = EventTerm(
+        func=mdp.randomize_rigid_body_mass,
+        mode="startup",
+        params={
+            "asset_cfg": SceneEntityCfg("robot", body_names="base"),
+            "mass_distribution_params": (-5.0, 5.0),
+            "operation": "add",
+        },
+    )
+
+    # reset
+    base_external_force_torque = EventTerm(
+        func=mdp.apply_external_force_torque,
+        mode="reset",
+        params={
+            "asset_cfg": SceneEntityCfg("robot", body_names="base"),
+            "force_range": (0.0, 0.0),
+            "torque_range": (-0.0, 0.0),
+        },
+    )
+
+    reset_base = EventTerm(
+        func=mdp.reset_root_state_uniform,
+        mode="reset",
+        params={
+            "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-3.14, 3.14)},
+            "velocity_range": {
+                "x": (-0.5, 0.5),
+                "y": (-0.5, 0.5),
+                "z": (-0.5, 0.5),
+                "roll": (-0.5, 0.5),
+                "pitch": (-0.5, 0.5),
+                "yaw": (-0.5, 0.5),
+            },
+        },
+    )
+
+    reset_robot_joints = EventTerm(
+        func=mdp.reset_joints_by_scale,
+        mode="reset",
+        params={
+            "position_range": (0.5, 1.5),
+            "velocity_range": (0.0, 0.0),
+        },
+    )
+
+    # interval
+    # comment for pit and gap
+    push_robot = EventTerm(
+        func=mdp.push_by_setting_velocity,
+        mode="interval",
+        interval_range_s=(3.0, 4.5),
+        params={"velocity_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5)}},
+    )
+
+
+@configclass
+class RewardsCfg:
+
+    # task rewards, eq. 1
+    task_reward = RewTerm(func=mdp.task_reward, weight=10.0, params={"reward_window": 1.0})
+    heading_reward = RewTerm(
+        func=mdp.heading_tracking, weight=10.0, params={"distance_threshold": 0.5, "reward_window": 1.0}
+    )
+
+    # penalties, eq. 2
+    joint_accel_l2 = RewTerm(func=mdp.joint_acc_l2, weight=-2.5e-7)
+    joint_torque_l2 = RewTerm(func=mdp.joint_torques_l2, weight=-1.0e-5)
+    undesired_contact = RewTerm(
+        func=mdp.undesired_contacts,
+        weight=-1.0,
+        params={
+            "sensor_cfg": SceneEntityCfg("contact_forces", body_names=".*THIGH"),
+            "threshold": 1.0,
+        },
+    )
+
+    action_rate_l2 = RewTerm(func=mdp.action_rate_l2, weight=-0.01)
+    feet_lin_acc_l2 = RewTerm(
+        func=mdp.feet_lin_acc_l2, weight=-4e-6, params={"robot_cfg": SceneEntityCfg("robot", body_names=".*FOOT")}
+    )
+    feet_rot_acc_l2 = RewTerm(
+        func=mdp.feet_rot_acc_l2, weight=-2e-7, params={"robot_cfg": SceneEntityCfg("robot", body_names=".*FOOT")}
+    )
+
+    illegal_contact_penalty = RewTerm(
+        func=mdp.illegal_contact_penalty,
+        weight=-3,
+        params={"sensor_cfg": SceneEntityCfg("contact_forces", body_names="base"), "threshold": 1.0},
+    )
+
+    # exploration eq. 3
+    exploration = RewTerm(func=mdp.exploration_reward, weight=1.0)
+
+    # stalling penalty eq. 4
+    stalling = RewTerm(func=mdp.stall_penalty, weight=-1.5, params={"distance_threshold": 0.2})
+
+
+@configclass
+class TerminationsCfg:
+    time_out = DoneTerm(func=mdp.time_out, time_out=True)
+
+    robot_drop = DoneTerm(
+        func=mdp.root_height_below_minimum,
+        params={
+            "minimum_height": -20,
+        },
+    )
+
+    base_contact = DoneTerm(
+        func=mdp.illegal_contact,
+        params={
+            "sensor_cfg": SceneEntityCfg("contact_forces", body_names="base"),
+            "threshold": 1.0,
+        },
+    )
+
+
+@configclass
+class CurriculumCfg:
+    terrain_levels = CurrTerm(func=mdp.terrain_levels_vel)  # type: ignore
+
+
+@configclass
+class AdvanceSkillsBaseEnvCfg(DataManagerBasedRLEnvCfg):
+    scene: SceneCfg = SceneCfg(num_envs=4096, env_spacing=10)
+    observations: ObservationsCfg = ObservationsCfg()
+    actions: ActionsCfg = ActionsCfg()
+    commands: CommandsCfg = CommandsCfg()
+    rewards: RewardsCfg = RewardsCfg()
+    terminations: TerminationsCfg = TerminationsCfg()
+    events: EventsCfg = EventsCfg()
+    curriculum: CurriculumCfg = CurriculumCfg()
+    viewer: ViewerCfg = ViewerCfg(eye=(1.0, 2.0, 2.0), origin_type="asset_body", asset_name="robot", body_name="base")
+
+    def __post_init__(self):
+        self.decimation = 4
+        self.episode_length_s = 6.0
+        self.sim.dt = 0.005
+        self.sim.render_interval = self.decimation
+        self.sim.disable_contact_processing = True
+        self.sim.physics_material = self.scene.terrain.physics_material
+        self.sim.physx.gpu_total_aggregate_pairs_capacity = 2**24
+        self.sim.physx.gpu_found_lost_pairs_capacity = 2**24
+        self.sim.physx.gpu_collision_stack_size = 2**27
+        self.sim.physx.gpu_max_rigid_patch_count = 5 * 2**16
+
+        # update sensor update periods
+        # we tick all the sensors based on the smallest update period (physics update period)
+        if self.scene.height_scanner is not None:
+            self.scene.height_scanner.update_period = self.decimation * self.sim.dt
+        if self.scene.contact_forces is not None:
+            self.scene.contact_forces.update_period = self.sim.dt
+
+        # check if terrain levels curriculum is enabled - if so, enable curriculum for terrain generator
+        # this generates terrains with increasing difficulty and is useful for training
+        if getattr(self.curriculum, "terrain_levels", None) is not None:
+            if self.scene.terrain.terrain_generator is not None:
+                self.scene.terrain.terrain_generator.curriculum = True
+        else:
+            if self.scene.terrain.terrain_generator is not None:
+                self.scene.terrain.terrain_generator.curriculum = False

source/uwlab_tasks/uwlab_tasks/manager_based/locomotion/advance_skills/advance_skills_env.py

@@ -0,0 +1,56 @@
+# Copyright (c) 2024-2025, The UW Lab Project Developers.
+# All Rights Reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from .advance_skills_base_env import AdvanceSkillsBaseEnvCfg
+from .terrains import EXTREME_STAIR, GAP, IRREGULAR_PILLAR_OBSTACLE, PIT, SLOPE_INV, SQUARE_PILLAR_OBSTACLE
+
+
+class GapEnvConfig(AdvanceSkillsBaseEnvCfg):
+    def __post_init__(self):
+        super().__post_init__()
+        self.scene.terrain.terrain_generator.sub_terrains = {"gap": GAP}
+
+
+class PitEnvConfig(AdvanceSkillsBaseEnvCfg):
+    def __post_init__(self):
+        super().__post_init__()
+        self.scene.terrain.terrain_generator.sub_terrains = {"pit": PIT}
+
+
+class ExtremeStairEnvConfig(AdvanceSkillsBaseEnvCfg):
+    def __post_init__(self):
+        super().__post_init__()
+        self.scene.terrain.terrain_generator.sub_terrains = {"pit": EXTREME_STAIR}
+
+
+class SlopeInvEnvConfig(AdvanceSkillsBaseEnvCfg):
+    def __post_init__(self):
+        super().__post_init__()
+        self.scene.terrain.terrain_generator.sub_terrains = {"slope_inv": SLOPE_INV}
+
+
+class SquarePillarObstacleEnvConfig(AdvanceSkillsBaseEnvCfg):
+    def __post_init__(self):
+        super().__post_init__()
+        self.scene.terrain.terrain_generator.sub_terrains = {"square_pillar_obstacle": SQUARE_PILLAR_OBSTACLE}
+
+
+class IrregularPillarObstacleEnvConfig(AdvanceSkillsBaseEnvCfg):
+    def __post_init__(self):
+        super().__post_init__()
+        self.scene.terrain.terrain_generator.sub_terrains = {"irregular_pillar_obstacle": IRREGULAR_PILLAR_OBSTACLE}
+
+
+class AdvanceSkillsEnvCfg(AdvanceSkillsBaseEnvCfg):
+    def __post_init__(self):
+        super().__post_init__()
+        self.scene.terrain.terrain_generator.sub_terrains = {
+            "gap": GAP,
+            "pit": PIT,
+            "extreme_stair": EXTREME_STAIR,
+            "slope_inv": SLOPE_INV,
+            "square_pillar_obstacle": SQUARE_PILLAR_OBSTACLE,
+            "irregular_pillar_obstacle": IRREGULAR_PILLAR_OBSTACLE,
+        }