support anymal d rough terrain

Author: ooctipus

source/isaaclab/isaaclab/envs/mdp/__init__.pyi

@@ -60,6 +60,7 @@ __all__ = [
     "randomize_rigid_body_material",
     "randomize_rigid_body_scale",
     "randomize_visual_color",
+    "randomize_visual_color_instanced",
     "randomize_visual_texture_material",
     "reset_joints_by_offset",
     "reset_joints_by_scale",
@@ -138,6 +139,8 @@ __all__ = [
     "joint_pos_out_of_manual_limit",
     "joint_vel_out_of_limit",
     "joint_vel_out_of_manual_limit",
+    "body_lin_vel_out_of_limit",
+    "body_ang_vel_out_of_limit",
     "root_height_below_minimum",
     "time_out",
 ]
@@ -201,6 +204,7 @@ from .events import (
     randomize_rigid_body_material,
     randomize_rigid_body_scale,
     randomize_visual_color,
+    randomize_visual_color_instanced,
     randomize_visual_texture_material,
     reset_joints_by_offset,
     reset_joints_by_scale,
@@ -287,6 +291,8 @@ from .terminations import (
     joint_pos_out_of_manual_limit,
     joint_vel_out_of_limit,
     joint_vel_out_of_manual_limit,
+    body_lin_vel_out_of_limit,
+    body_ang_vel_out_of_limit,
     root_height_below_minimum,
     time_out,
 )

source/isaaclab/isaaclab/envs/mdp/terminations.py

@@ -150,6 +150,41 @@ def joint_effort_out_of_limit(
     return torch.any(out_of_limits, dim=1)
 
 
+"""
+Body velocity terminations.
+"""
+
+
+def body_lin_vel_out_of_limit(
+    env: ManagerBasedRLEnv,
+    max_velocity: float,
+    asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
+) -> torch.Tensor:
+    """Terminate when any body's linear velocity magnitude exceeds the limit [m/s].
+
+    Also terminates on NaN velocities, which indicate a solver singularity.
+    """
+    asset: Articulation = env.scene[asset_cfg.name]
+    body_vel = wp.to_torch(asset.data.body_lin_vel_w)[:, asset_cfg.body_ids]
+    speed = torch.linalg.norm(body_vel, dim=-1)
+    return torch.any((speed > max_velocity) | torch.isnan(speed), dim=1)
+
+
+def body_ang_vel_out_of_limit(
+    env: ManagerBasedRLEnv,
+    max_velocity: float,
+    asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
+) -> torch.Tensor:
+    """Terminate when any body's angular velocity magnitude exceeds the limit [rad/s].
+
+    Also terminates on NaN velocities, which indicate a solver singularity.
+    """
+    asset: Articulation = env.scene[asset_cfg.name]
+    body_vel = wp.to_torch(asset.data.body_ang_vel_w)[:, asset_cfg.body_ids]
+    speed = torch.linalg.norm(body_vel, dim=-1)
+    return torch.any((speed > max_velocity) | torch.isnan(speed), dim=1)
+
+
 """
 Contact sensor.
 """

source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/config/anymal_d/rough_env_cfg.py

@@ -3,13 +3,13 @@
 #
 # SPDX-License-Identifier: BSD-3-Clause
 
+from isaaclab_newton.physics import MJWarpSolverCfg, NewtonCfg, NewtonCollisionPipelineCfg
+from isaaclab_physx.physics import PhysxCfg
+
+from isaaclab.sim import SimulationCfg
 from isaaclab.utils import configclass
 
-from isaaclab_tasks.manager_based.locomotion.velocity.velocity_env_cfg import (
-    EventsCfg,
-    LocomotionVelocityRoughEnvCfg,
-    StartupEventsCfg,
-)
+from isaaclab_tasks.manager_based.locomotion.velocity.velocity_env_cfg import LocomotionVelocityRoughEnvCfg
 from isaaclab_tasks.utils import PresetCfg
 
 ##
@@ -19,28 +19,33 @@
 
 
 @configclass
-class AnymalDPhysxEventsCfg(EventsCfg, StartupEventsCfg):
-    pass
-
-
-@configclass
-class AnymalDEventsCfg(PresetCfg):
-    default = AnymalDPhysxEventsCfg()
-    newton = EventsCfg()
+class RoughPhysicsCfg(PresetCfg):
+    default = PhysxCfg(gpu_max_rigid_patch_count=10 * 2**15)
+    newton = NewtonCfg(
+        solver_cfg=MJWarpSolverCfg(
+            njmax=200,
+            nconmax=100,
+            cone="pyramidal",
+            impratio=1.0,
+            integrator="implicitfast",
+            use_mujoco_contacts=False,
+        ),
+        collision_cfg=NewtonCollisionPipelineCfg(max_triangle_pairs=2_500_000),
+        num_substeps=1,
+        debug_mode=False,
+    )
     physx = default
 
-
 @configclass
 class AnymalDRoughEnvCfg(LocomotionVelocityRoughEnvCfg):
-    events: AnymalDEventsCfg = AnymalDEventsCfg()
+    sim: SimulationCfg = SimulationCfg(physics=RoughPhysicsCfg())
 
     def __post_init__(self):
         # post init of parent
         super().__post_init__()
         # switch robot to anymal-d
         self.scene.robot = ANYMAL_D_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")
 
-
 @configclass
 class AnymalDRoughEnvCfg_PLAY(AnymalDRoughEnvCfg):
     def __post_init__(self):

source/isaaclab_tasks/isaaclab_tasks/manager_based/locomotion/velocity/velocity_env_cfg.py

@@ -26,8 +26,9 @@
 from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR, ISAACLAB_NUCLEUS_DIR
 from isaaclab.utils.noise import UniformNoiseCfg as Unoise
 
+
 import isaaclab_tasks.manager_based.locomotion.velocity.mdp as mdp
-from isaaclab_tasks.utils import PresetCfg
+from isaaclab_tasks.utils import PresetCfg, preset
 
 ##
 # Pre-defined configs
@@ -161,6 +162,49 @@ def __post_init__(self):
 @configclass
 class EventsCfg:
     """Configuration for events."""
+    # startup
+    physics_material = EventTerm(
+        func=mdp.randomize_rigid_body_material,
+        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",
+        },
+    )
+
+    base_com = preset(
+        default=EventTerm(
+            func=mdp.randomize_rigid_body_com,
+            mode="startup",
+            params={
+                "asset_cfg": SceneEntityCfg("robot", body_names="base"),
+                "com_range": {"x": (-0.05, 0.05), "y": (-0.05, 0.05), "z": (-0.01, 0.01)},
+            },
+        ),
+        newton=None
+    )
+
+    collider_offsets = EventTerm(
+        func=mdp.randomize_rigid_body_collider_offsets,
+        mode="startup",
+        params={
+            "asset_cfg": SceneEntityCfg("robot", body_names=".*"),
+            "contact_offset_distribution_params": (0.01, 0.01)
+        },
+    )
 
     # reset
     base_external_force_torque = EventTerm(
@@ -207,40 +251,6 @@ class EventsCfg:
     )
 
 
-@configclass
-class StartupEventsCfg:
-    # startup
-    physics_material = EventTerm(
-        func=mdp.randomize_rigid_body_material,
-        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",
-        },
-    )
-
-    base_com = EventTerm(
-        func=mdp.randomize_rigid_body_com,
-        mode="startup",
-        params={
-            "asset_cfg": SceneEntityCfg("robot", body_names="base"),
-            "com_range": {"x": (-0.05, 0.05), "y": (-0.05, 0.05), "z": (-0.01, 0.01)},
-        },
-    )
-
 
 @configclass
 class RewardsCfg:
@@ -287,6 +297,10 @@ class TerminationsCfg:
         func=mdp.illegal_contact,
         params={"sensor_cfg": SceneEntityCfg("contact_forces", body_names="base"), "threshold": 1.0},
     )
+    body_lin_vel = DoneTerm(
+        func=mdp.body_lin_vel_out_of_limit,
+        params={"max_velocity": 20.0, "asset_cfg": SceneEntityCfg("robot", body_names=".*")},
+    )
 
 
 @configclass
@@ -314,7 +328,7 @@ class LocomotionVelocityRoughEnvCfg(ManagerBasedRLEnvCfg):
     # MDP settings
     rewards: RewardsCfg = RewardsCfg()
     terminations: TerminationsCfg = TerminationsCfg()
-    events: EventsCfg = MISSING
+    events: EventsCfg = EventsCfg()
     curriculum: CurriculumCfg = CurriculumCfg()
 
     def __post_init__(self):