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 # ---------------------------------------------------------------------------