From 9f95f70efba082d2ccc7d94d3e58f68bd0126e99 Mon Sep 17 00:00:00 2001 From: Henry Hu Date: Wed, 30 Sep 2026 04:25:53 +0800 Subject: [PATCH] Counterbalance the biped air-time reward on Cassie/Go2/H1 (#7813) # Description Cassie/H1/Go2 velocity tasks converge on a badly lopsided gait: one leg stays planted while the other hangs, which reads as "standing on one leg" in playback. `feet_air_time_positive_biped` scores the length of the current single-stance phase and takes the minimum over the two feet: ```python in_mode_time = torch.where(in_contact, contact_time, air_time) single_stance = torch.sum(in_contact.int(), dim=1) == 1 reward = torch.min(torch.where(single_stance.unsqueeze(-1), in_mode_time, 0.0), dim=1)[0] reward = torch.clamp(reward, max=threshold) ``` Nothing in it distinguishes left from right or asks the feet to alternate. Holding one foot planted and the other airborne lets both timers run and sits at the clamp, while a symmetric gait scores lower. This PR adds `feet_air_time_variance`, the variance across feet of the last completed swing and stance durations, and enables it on all the three tasks. It vanishes when both feet are loaded equally, so it prices the exploit without constraining anything else. ## Type of change - New feature (non-breaking change which adds functionality) Policies trained on these two tasks will differ from ones trained before this change. Setting `env.rewards.air_time_variance.weight=0.0` recovers the previous reward exactly. ## Release backport - [ ] Backport this pull request to the active release branch after it merges into `develop` ## Checklist - [x] I have read and understood the [contribution guidelines](https://isaac-sim.github.io/IsaacLab/main/source/refs/contributing.html) - [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with `./isaaclab.sh --format` - [x] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [x] I have added a changelog fragment under `source//changelog.d/` for every touched package - [x] I have added my name to the `CONTRIBUTORS.md` or my name already exists there --------- Co-authored-by: ooctipus --- .../henry-go2-base-height-inf-guard.rst | 7 ++ source/isaaclab/isaaclab/envs/mdp/rewards.py | 14 +++- .../changelog.d/warp-base-height-reward.rst | 5 ++ .../isaaclab_experimental/envs/mdp/rewards.py | 37 ++++++++++ .../test/envs/mdp/test_capture_safety.py | 71 +++++++++++++++++++ .../henry-cassie-gait-symmetry.minor.rst | 13 ++++ .../henry-go2-crouch-and-gait-fix.minor.rst | 12 ++++ .../velocity/config/cassie/flat_env_cfg.py | 1 + .../velocity/config/cassie/rough_env_cfg.py | 9 +++ .../core/velocity/config/go2/flat_env_cfg.py | 1 + .../core/velocity/config/go2/rough_env_cfg.py | 33 ++++++++- .../core/velocity/config/h1/flat_env_cfg.py | 23 +++++- .../core/velocity/mdp/__init__.pyi | 2 + .../core/velocity/mdp/rewards.py | 33 +++++++++ .../warp-air-time-variance-reward.rst | 5 ++ .../core/velocity/mdp/__init__.pyi | 2 + .../core/velocity/mdp/rewards.py | 50 +++++++++++++ 17 files changed, 314 insertions(+), 4 deletions(-) create mode 100644 source/isaaclab/changelog.d/henry-go2-base-height-inf-guard.rst create mode 100644 source/isaaclab_experimental/changelog.d/warp-base-height-reward.rst create mode 100644 source/isaaclab_tasks/changelog.d/henry-cassie-gait-symmetry.minor.rst create mode 100644 source/isaaclab_tasks/changelog.d/henry-go2-crouch-and-gait-fix.minor.rst create mode 100644 source/isaaclab_tasks_experimental/changelog.d/warp-air-time-variance-reward.rst diff --git a/source/isaaclab/changelog.d/henry-go2-base-height-inf-guard.rst b/source/isaaclab/changelog.d/henry-go2-base-height-inf-guard.rst new file mode 100644 index 000000000000..a64094f7ef1f --- /dev/null +++ b/source/isaaclab/changelog.d/henry-go2-base-height-inf-guard.rst @@ -0,0 +1,7 @@ +Fixed +^^^^^ + +* Fixed :func:`~isaaclab.envs.mdp.rewards.base_height_l2` returning an infinite reward for any + environment whose own height scan contained a missed ray, since ``ray_hits_w`` reports ``inf`` + for a ray that finds nothing within its ``max_distance`` rather than a clamped value. Now only + finite ray hits are averaged into the terrain-adjusted target height. diff --git a/source/isaaclab/isaaclab/envs/mdp/rewards.py b/source/isaaclab/isaaclab/envs/mdp/rewards.py index 0c96e700f76d..127c8cd3d9b4 100644 --- a/source/isaaclab/isaaclab/envs/mdp/rewards.py +++ b/source/isaaclab/isaaclab/envs/mdp/rewards.py @@ -131,12 +131,22 @@ def base_height_l2( Note: For flat terrain, target height is in the world frame. For rough terrain, sensor readings can adjust the target height to account for the terrain. + + Note: + A ray that hits nothing within its ``max_distance`` reports ``inf`` rather than a clamped + value, so only finite hits are averaged into the terrain-adjusted target height. An + environment whose rays all miss falls back to the unadjusted ``target_height``. """ asset: RigidObject = env.scene[asset_cfg.name] if sensor_cfg is not None: sensor: RayCaster = env.scene[sensor_cfg.name] - # Adjust the target height using the sensor data - adjusted_target_height = target_height + torch.mean(sensor.data.ray_hits_w.torch[..., 2], dim=1) + # Adjust the target height using the sensor data, ignoring rays that missed the terrain + ray_hits_z = sensor.data.ray_hits_w.torch[..., 2] + finite_hits = torch.isfinite(ray_hits_z) + safe_hits = torch.where(finite_hits, ray_hits_z, torch.zeros_like(ray_hits_z)) + hit_count = finite_hits.sum(dim=1).clamp(min=1) + terrain_height = safe_hits.sum(dim=1) / hit_count + adjusted_target_height = target_height + terrain_height else: # Use the provided target height directly for flat terrain adjusted_target_height = target_height diff --git a/source/isaaclab_experimental/changelog.d/warp-base-height-reward.rst b/source/isaaclab_experimental/changelog.d/warp-base-height-reward.rst new file mode 100644 index 000000000000..e90c5e8cc9ef --- /dev/null +++ b/source/isaaclab_experimental/changelog.d/warp-base-height-reward.rst @@ -0,0 +1,5 @@ +Fixed +^^^^^ + +* Added the Warp implementation of ``base_height_l2``, restoring Warp frontend support for + Unitree Go2 tasks with the base-height reward enabled. diff --git a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py index bcebee08725a..880697b562f2 100644 --- a/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py +++ b/source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py @@ -206,6 +206,43 @@ def flat_orientation_l2(env: ManagerBasedRLEnv, out, asset_cfg: SceneEntityCfg = ) +@wp.kernel +def _base_height_l2_kernel( + root_pose_w: wp.array(dtype=wp.transformf), + ray_hits_w: wp.array(dtype=wp.vec3f, ndim=2), + target_height: float, + out: wp.array(dtype=wp.float32), +): + i = wp.tid() + height_sum = float(0.0) + hit_count = int(0) + if ray_hits_w: + for ray in range(ray_hits_w.shape[1]): + height = ray_hits_w[i, ray][2] + if wp.isfinite(height): + height_sum += height + hit_count += 1 + error = wp.transform_get_translation(root_pose_w[i])[2] - target_height - height_sum / float(wp.max(hit_count, 1)) + out[i] = error * error + + +def base_height_l2( + env: ManagerBasedRLEnv, + out: wp.array(dtype=wp.float32), + target_height: float, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + sensor_cfg: SceneEntityCfg | None = None, +) -> None: + """Penalize squared base-height error from ``target_height`` [m], relative to finite terrain hits.""" + ray_hits = env.scene[sensor_cfg.name].data.ray_hits_w if sensor_cfg is not None else None + wp.launch( + _base_height_l2_kernel, + dim=env.num_envs, + inputs=[env.scene[asset_cfg.name].data.root_link_pose_w, ray_hits, target_height, out], + device=env.device, + ) + + """ Joint penalties. """ diff --git a/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py b/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py index 7fbc95efebe7..ea69b3308fc4 100644 --- a/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py +++ b/source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py @@ -30,6 +30,7 @@ import inspect import pkgutil from collections.abc import Callable +from functools import partial from types import SimpleNamespace import numpy as np @@ -44,25 +45,32 @@ import isaaclab_experimental.envs.mdp.terminations as warp_term import isaaclab_tasks_experimental.core as tasks_experimental_core import isaaclab_tasks_experimental.core.locomotion.mdp.rewards as warp_loco_rew +import isaaclab_tasks_experimental.core.velocity.mdp.rewards as warp_velocity_rew from parity_helpers import ( DEVICE, NUM_BODIES, NUM_ENVS, MockArticulation, MockArticulationData, + MockCommandManager, MockPoseCommandManager, MockScene, + MockSensorCfg, MockTerminationManager, assert_close, assert_equal, + copy_np_to_wp, make_pose_command_term, mutate_body_data, + proxy_array, ) import isaaclab.envs.mdp.rewards as stable_rew import isaaclab.envs.mdp.terminations as stable_term from isaaclab.managers.manager_term_cfg import RewardTermCfg, TerminationTermCfg +import isaaclab_tasks.core.velocity.mdp.rewards as stable_velocity_rew + @dataclasses.dataclass(frozen=True) class CaptureCase: @@ -236,6 +244,53 @@ def mutate() -> None: ) +def _build_base_height_l2(terrain: bool) -> CaptureCase: + art_data = MockArticulationData(num_bodies=NUM_BODIES) + hits = np.random.default_rng(17).normal(size=(NUM_ENVS, 4, 3)).astype(np.float32) + hits[::2, 0, 2] = np.inf + hits[::3, 1, 2] = -np.inf + hits[::5, :, 2] = np.nan # No finite hits: use the flat-ground target. + ray_hits = proxy_array(hits, dtype=wp.vec3f, device=DEVICE) + scene = { + "robot": MockArticulation(art_data), + "height_scanner": SimpleNamespace(data=SimpleNamespace(ray_hits_w=ray_hits)), + } + env = SimpleNamespace(scene=scene, num_envs=NUM_ENVS, device=DEVICE) + params = dict(target_height=0.35, sensor_cfg=MockSensorCfg("height_scanner") if terrain else None) + + def mutate(): + mutate_body_data(art_data) + copy_np_to_wp(ray_hits, hits + 0.3) + + return CaptureCase(warp_rew.base_height_l2, stable_rew.base_height_l2, env, env, params, mutate) + + +def _build_feet_air_time_variance(body_ids=(0, 2)) -> CaptureCase: + rng = np.random.default_rng(23) + air = rng.uniform(0.0, 1.0, (NUM_ENVS, NUM_BODIES)).astype(np.float32) + contact = rng.uniform(0.0, 1.0, air.shape).astype(np.float32) + data = SimpleNamespace( + last_air_time=proxy_array(air, device=DEVICE), last_contact_time=proxy_array(contact, device=DEVICE) + ) + scene = SimpleNamespace(sensors={"contact_sensor": SimpleNamespace(data=data)}) + command = torch.tensor(rng.normal(size=(NUM_ENVS, 3)), dtype=torch.float32, device=DEVICE) + command[::2, :2] = 0.0 # Yaw-only commands must not incur this penalty. + command[1, :2] = torch.tensor([0.06, 0.07], device=DEVICE) + env = SimpleNamespace( + scene=scene, command_manager=MockCommandManager(command, None), num_envs=NUM_ENVS, device=DEVICE + ) + params = dict(command_name="vel", sensor_cfg=MockSensorCfg(body_ids=list(body_ids)), max_time=0.5) + + def mutate(): + copy_np_to_wp(data.last_air_time, 1.0 - air) + copy_np_to_wp(data.last_contact_time, 1.0 - contact) + command.copy_(command.roll(1, dims=0)) + + return CaptureCase( + warp_velocity_rew.feet_air_time_variance, stable_velocity_rew.feet_air_time_variance, env, env, params, mutate + ) + + _SHARED_REW = "isaaclab_experimental.envs.mdp.rewards" _SHARED_TERM = "isaaclab_experimental.envs.mdp.terminations" _LOCO_REW = "isaaclab_tasks_experimental.core.locomotion.mdp.rewards" @@ -245,6 +300,14 @@ def mutate() -> None: CaptureSpec(_SHARED_TERM, "pose_command_success", "termination", _build_pose_command_success), CaptureSpec(_LOCO_REW, "terminated_penalty", "reward", _build_terminated_penalty), CaptureSpec(_LOCO_REW, "survival_success_rate", "reward", _build_survival_success_rate, expect_nonzero=False), + CaptureSpec(_SHARED_REW, "base_height_l2", "reward", partial(_build_base_height_l2, terrain=False)), + CaptureSpec(_SHARED_REW, "base_height_l2", "reward", partial(_build_base_height_l2, terrain=True)), + CaptureSpec( + "isaaclab_tasks_experimental.core.velocity.mdp.rewards", + "feet_air_time_variance", + "reward", + _build_feet_air_time_variance, + ), ] @@ -458,6 +521,14 @@ def test_term_is_capture_safe(spec: CaptureSpec): assert expected.any(), "mutated inputs produced a degenerate expectation; the replay proves little" +def test_feet_air_time_variance_requires_two_feet(): + case = _build_feet_air_time_variance(body_ids=[0]) + with pytest.raises(RuntimeError, match="at least two bodies"): + case.stable_fn(case.stable_env, **case.params) + with pytest.raises(RuntimeError, match="at least two bodies"): + case.warp_fn(case.warp_env, wp.empty(NUM_ENVS, dtype=wp.float32, device=DEVICE), **case.params) + + def test_survival_success_rate_reset_is_capture_safe(): """The metric is produced by ``reset()``, not ``__call__``, so capture it separately. diff --git a/source/isaaclab_tasks/changelog.d/henry-cassie-gait-symmetry.minor.rst b/source/isaaclab_tasks/changelog.d/henry-cassie-gait-symmetry.minor.rst new file mode 100644 index 000000000000..0da0e419b2b7 --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/henry-cassie-gait-symmetry.minor.rst @@ -0,0 +1,13 @@ +Added +^^^^^ + +* Added :func:`~isaaclab_tasks.core.velocity.mdp.rewards.feet_air_time_variance`, penalizing an + uneven swing/stance split between a biped's feet. + +Changed +^^^^^^^ + +* Changed ``Isaac-Velocity-Flat-Cassie``, ``Isaac-Velocity-Rough-Cassie`` and + ``Isaac-Velocity-Flat-H1`` to include the new ``air_time_variance`` reward term, which fixes the + tilted gait these tasks otherwise converge on. Policies trained on them will differ from ones + trained before this change; set the weight to ``0.0`` to recover the previous reward. diff --git a/source/isaaclab_tasks/changelog.d/henry-go2-crouch-and-gait-fix.minor.rst b/source/isaaclab_tasks/changelog.d/henry-go2-crouch-and-gait-fix.minor.rst new file mode 100644 index 000000000000..022ede555879 --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/henry-go2-crouch-and-gait-fix.minor.rst @@ -0,0 +1,12 @@ +Changed +^^^^^^^ + +* Changed ``Isaac-Velocity-Flat-UnitreeGo2`` and ``Isaac-Velocity-Rough-UnitreeGo2`` to add a + ``base_height_l2`` reward term at weight ``-30.0`` (rough, height-scanner adjusted) and disabled + on flat via ``sensor_cfg=None``. Without it, Go2's leg colliders let a policy rest its weight on + the legs and collect the full episode-length alive bonus from a permanent crouch, since torso + contact never fires and flat-orientation reward can't see it. +* Changed ``action_rate_l2`` on both tasks from the shared default of ``-0.01`` to ``-0.005``. The + stronger penalty locked one hind foot into a low-amplitude, dragging gait on flat terrain with + the Newton backend. Policies trained on these tasks will differ from ones trained before this + change. diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/cassie/flat_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/cassie/flat_env_cfg.py index 73804d550de0..9112a0757808 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/cassie/flat_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/cassie/flat_env_cfg.py @@ -31,6 +31,7 @@ def __post_init__(self): # rewards self.rewards.flat_orientation_l2.weight = -2.5 self.rewards.feet_air_time.weight = 5.0 + self.rewards.air_time_variance.weight = -10.0 self.rewards.joint_deviation_hip.params["asset_cfg"].joint_names = ["hip_rotation_.*"] # curriculum self.curriculum.terrain_levels = None diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/cassie/rough_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/cassie/rough_env_cfg.py index 84f38bed9ffe..1c4b565934bb 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/cassie/rough_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/cassie/rough_env_cfg.py @@ -42,6 +42,15 @@ class CassieRewardsCfg(RewardsCfg): weight=-0.2, params={"asset_cfg": SceneEntityCfg("robot", joint_names=["toe_joint_.*"])}, ) + # fixes the tilted gait + air_time_variance = RewTerm( + func=mdp.feet_air_time_variance, + weight=-5.0, + params={ + "sensor_cfg": SceneEntityCfg("contact_forces", body_names=".*toe"), + "command_name": "base_velocity", + }, + ) # penalize toe joint limits dof_pos_limits = RewTerm( func=mdp.joint_pos_limits, diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/go2/flat_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/go2/flat_env_cfg.py index 46f3f9effadc..4237f735d0d0 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/go2/flat_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/go2/flat_env_cfg.py @@ -31,5 +31,6 @@ def __post_init__(self): # rewards self.rewards.flat_orientation_l2.weight = -2.5 self.rewards.feet_air_time.weight = 0.25 + self.rewards.base_height_l2.params["sensor_cfg"] = None # curriculum self.curriculum.terrain_levels = None diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/go2/rough_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/go2/rough_env_cfg.py index bcd1c051417c..71160e332795 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/go2/rough_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/go2/rough_env_cfg.py @@ -5,17 +5,42 @@ """Configuration for the Unitree Go2 velocity-tracking environment on rough terrain.""" +from isaaclab.managers import RewardTermCfg as RewTerm +from isaaclab.managers import SceneEntityCfg from isaaclab.utils import configclass, replace from isaaclab_assets.robots.unitree import UNITREE_GO2_CFG -from ...velocity_env_cfg import LocomotionVelocityRoughEnvCfg +from ... import mdp +from ...velocity_env_cfg import LocomotionVelocityRoughEnvCfg, RewardsCfg + +# UNITREE_GO2_CFG.init_state.pos is (0, 0, 0.4) -- the asset's own authored standing height. +_GO2_STANDING_HEIGHT = 0.4 + + +@configclass +class UnitreeGo2Rewards(RewardsCfg): + """Adds a base-height term on top of the shared velocity rewards. + + The stock velocity task has no term constraining base height. Go2 carries colliders on its + legs (not just the feet), so a policy can rest its weight on the legs and collect the full + episode-length alive bonus from a permanent crouch -- torso-contact termination never fires + and flat-orientation reward can't see it, since the torso stays level while the robot sinks. + """ + + base_height_l2 = RewTerm( + func=mdp.base_height_l2, + weight=0.0, + params={"target_height": _GO2_STANDING_HEIGHT}, + ) @configclass class UnitreeGo2RoughEnvCfg(LocomotionVelocityRoughEnvCfg): """Configuration for the Unitree Go2 velocity-tracking environment on rough terrain.""" + rewards: UnitreeGo2Rewards = UnitreeGo2Rewards() + def __post_init__(self): super().__post_init__() @@ -42,5 +67,11 @@ def __post_init__(self): self.rewards.track_lin_vel_xy_exp.weight = 1.5 self.rewards.track_ang_vel_z_exp.weight = 0.75 self.rewards.dof_acc_l2.weight = -2.5e-7 + # base_height_l2: rough terrain, adjust target height using the height-scanner rays. + self.rewards.base_height_l2.weight = -30.0 + self.rewards.base_height_l2.params["sensor_cfg"] = SceneEntityCfg("height_scanner") + # action_rate_l2: halved from the shared default of -0.01. The stronger penalty locked + # one hind foot into a low-amplitude, dragging gait on flat terrain with Newton. + self.rewards.action_rate_l2.weight = -0.005 # terminations self.terminations.base_contact.params["sensor_cfg"].body_names = "base" diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/h1/flat_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/h1/flat_env_cfg.py index 1db600e46532..2d98123dff34 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/h1/flat_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/config/h1/flat_env_cfg.py @@ -5,15 +5,36 @@ """Configuration for the Unitree H1 velocity-tracking environment on flat terrain.""" +from isaaclab.managers import RewardTermCfg as RewTerm +from isaaclab.managers import SceneEntityCfg from isaaclab.utils import configclass -from .rough_env_cfg import H1RoughEnvCfg +import isaaclab_tasks.core.velocity.mdp as mdp + +from .rough_env_cfg import H1Rewards, H1RoughEnvCfg + + +@configclass +class H1FlatRewards(H1Rewards): + """Reward terms for the MDP.""" + + # fixes the tilted gait + air_time_variance = RewTerm( + func=mdp.feet_air_time_variance, + weight=-2.0, + params={ + "sensor_cfg": SceneEntityCfg("contact_forces", body_names=".*ankle_link"), + "command_name": "base_velocity", + }, + ) @configclass class H1FlatEnvCfg(H1RoughEnvCfg): """Configuration for the Unitree H1 velocity-tracking environment on flat terrain.""" + rewards: H1FlatRewards = H1FlatRewards() + def __post_init__(self): super().__post_init__() diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/mdp/__init__.pyi b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/mdp/__init__.pyi index 025fecca866a..88a90656769f 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/mdp/__init__.pyi +++ b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/mdp/__init__.pyi @@ -6,6 +6,7 @@ __all__ = [ "feet_air_time", "feet_air_time_positive_biped", + "feet_air_time_variance", "feet_slide", "stand_still_joint_deviation_l1", "terrain_levels_vel", @@ -18,6 +19,7 @@ from .curriculums import terrain_levels_vel from .rewards import ( feet_air_time, feet_air_time_positive_biped, + feet_air_time_variance, feet_slide, stand_still_joint_deviation_l1, track_ang_vel_z_world_exp, diff --git a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/mdp/rewards.py index 45a9a6b42125..83efe1ebe697 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/core/velocity/mdp/rewards.py +++ b/source/isaaclab_tasks/isaaclab_tasks/core/velocity/mdp/rewards.py @@ -67,6 +67,39 @@ def feet_air_time_positive_biped( return reward +def feet_air_time_variance( + env: ManagerBasedRLEnv, command_name: str, sensor_cfg: SceneEntityCfg, max_time: float = 0.5 +) -> torch.Tensor: + """Penalize an uneven swing/stance split between the feet. + + The penalty is the variance across feet of the last completed swing and stance durations [s], + each clipped at ``max_time`` [s], and is zero for near-zero commands. + + Raises: + RuntimeError: If ``sensor_cfg`` resolves to fewer than two bodies, where :func:`torch.var` + would return ``nan``. + """ + # extract the used quantities (to enable type-hinting) + contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] + # compute the penalty + last_air_time = contact_sensor.data.last_air_time.torch[:, sensor_cfg.body_ids] + last_contact_time = contact_sensor.data.last_contact_time.torch[:, sensor_cfg.body_ids] + # ``body_ids`` may be a slice, so count the bodies the selection actually resolved to + num_feet = last_air_time.shape[1] + if num_feet < 2: + raise RuntimeError( + f"'feet_air_time_variance' needs at least two bodies to take a variance across feet, but" + f" '{sensor_cfg.name}' resolved to {num_feet}. Widen 'body_names' on its 'sensor_cfg' to" + " cover every foot." + ) + penalty = torch.var(torch.clip(last_air_time, max=max_time), dim=1) + torch.var( + torch.clip(last_contact_time, max=max_time), dim=1 + ) + # no penalty for zero command + penalty *= torch.linalg.norm(env.command_manager.get_command(command_name)[:, :2], dim=1) > 0.1 + return penalty + + def feet_slide( env: ManagerBasedRLEnv, sensor_cfg: SceneEntityCfg, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") ) -> torch.Tensor: diff --git a/source/isaaclab_tasks_experimental/changelog.d/warp-air-time-variance-reward.rst b/source/isaaclab_tasks_experimental/changelog.d/warp-air-time-variance-reward.rst new file mode 100644 index 000000000000..646b97634a7a --- /dev/null +++ b/source/isaaclab_tasks_experimental/changelog.d/warp-air-time-variance-reward.rst @@ -0,0 +1,5 @@ +Fixed +^^^^^ + +* Added the Warp implementation of ``feet_air_time_variance``, restoring Warp frontend support + for Cassie and H1 tasks with the gait-symmetry reward enabled. diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/__init__.pyi b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/__init__.pyi index ae4b6d7af547..c9a51b09123b 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/__init__.pyi +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/__init__.pyi @@ -7,6 +7,7 @@ __all__ = [ "terrain_levels_vel", "feet_air_time", "feet_air_time_positive_biped", + "feet_air_time_variance", "feet_slide", "stand_still_joint_deviation_l1", "track_ang_vel_z_world_exp", @@ -22,6 +23,7 @@ from .curriculums import terrain_levels_vel from .rewards import ( feet_air_time, feet_air_time_positive_biped, + feet_air_time_variance, feet_slide, stand_still_joint_deviation_l1, track_ang_vel_z_world_exp, diff --git a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/rewards.py b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/rewards.py index 345e8844b7fc..4861b58e5e33 100644 --- a/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/rewards.py +++ b/source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/velocity/mdp/rewards.py @@ -138,6 +138,56 @@ def feet_air_time_positive_biped(env, out, command_name: str, threshold: float, ) +@wp.kernel +def _feet_air_time_variance_kernel( + air_time: wp.array(dtype=wp.float32, ndim=2), + contact_time: wp.array(dtype=wp.float32, ndim=2), + body_ids: wp.array(dtype=wp.int32), + command: wp.array(dtype=wp.float32, ndim=2), + max_time: float, + out: wp.array(dtype=wp.float32), +): + i = wp.tid() + mean = wp.vec2f() + for foot in range(body_ids.shape[0]): + body = body_ids[foot] + mean += wp.vec2f(wp.min(air_time[i, body], max_time), wp.min(contact_time[i, body], max_time)) + mean /= float(body_ids.shape[0]) + variance = float(0.0) + for foot in range(body_ids.shape[0]): + body = body_ids[foot] + times = wp.vec2f(wp.min(air_time[i, body], max_time), wp.min(contact_time[i, body], max_time)) + variance += wp.length_sq(times - mean) + moving = wp.length(wp.vec2f(command[i, 0], command[i, 1])) > 0.1 + out[i] = wp.where(moving, variance / float(body_ids.shape[0] - 1), 0.0) + + +def feet_air_time_variance( + env: ManagerBasedRLEnv, + out: wp.array(dtype=wp.float32), + command_name: str, + sensor_cfg: SceneEntityCfg, + max_time: float = 0.5, +) -> None: + """Penalize sample variance of foot air/contact durations, capped at ``max_time`` [s], while moving.""" + if sensor_cfg.body_ids_wp.shape[0] < 2: + raise RuntimeError("'feet_air_time_variance' needs at least two bodies to take a variance across feet.") + data = env.scene.sensors[sensor_cfg.name].data + wp.launch( + _feet_air_time_variance_kernel, + dim=env.num_envs, + inputs=[ + data.last_air_time, + data.last_contact_time, + sensor_cfg.body_ids_wp, + wp.from_torch(env.command_manager.get_command(command_name)), + max_time, + out, + ], + device=env.device, + ) + + # --------------------------------------------------------------------------- # feet_slide # ---------------------------------------------------------------------------