|
18 | 18 | from omni.isaac.lab.managers import SceneEntityCfg
|
19 | 19 | from omni.isaac.lab.managers.manager_base import ManagerTermBase
|
20 | 20 | from omni.isaac.lab.managers.manager_term_cfg import RewardTermCfg
|
21 |
| -from omni.isaac.lab.sensors import ContactSensor |
| 21 | +from omni.isaac.lab.sensors import ContactSensor, RayCaster |
22 | 22 |
|
23 | 23 | if TYPE_CHECKING:
|
24 | 24 | from omni.isaac.lab.envs import ManagerBasedRLEnv
|
@@ -98,17 +98,28 @@ def flat_orientation_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = Scen
|
98 | 98 |
|
99 | 99 |
|
100 | 100 | def base_height_l2(
|
101 |
| - env: ManagerBasedRLEnv, target_height: float, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") |
| 101 | + env: ManagerBasedRLEnv, |
| 102 | + target_height: float, |
| 103 | + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), |
| 104 | + sensor_cfg: SceneEntityCfg | None = None, |
102 | 105 | ) -> torch.Tensor:
|
103 | 106 | """Penalize asset height from its target using L2 squared kernel.
|
104 | 107 |
|
105 | 108 | Note:
|
106 |
| - Currently, it assumes a flat terrain, i.e. the target height is in the world frame. |
| 109 | + For flat terrain, target height is in the world frame. For rough terrain, |
| 110 | + sensor readings can adjust the target height to account for the terrain. |
107 | 111 | """
|
108 | 112 | # extract the used quantities (to enable type-hinting)
|
109 | 113 | asset: RigidObject = env.scene[asset_cfg.name]
|
110 |
| - # TODO: Fix this for rough-terrain. |
111 |
| - return torch.square(asset.data.root_pos_w[:, 2] - target_height) |
| 114 | + if sensor_cfg is not None: |
| 115 | + sensor: RayCaster = env.scene[sensor_cfg.name] |
| 116 | + # Adjust the target height using the sensor data |
| 117 | + adjusted_target_height = target_height + sensor.data.pos_w[:, 2] |
| 118 | + else: |
| 119 | + # Use the provided target height directly for flat terrain |
| 120 | + adjusted_target_height = target_height |
| 121 | + # Compute the L2 squared penalty |
| 122 | + return torch.square(asset.data.root_pos_w[:, 2] - adjusted_target_height) |
112 | 123 |
|
113 | 124 |
|
114 | 125 | def body_lin_acc_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
|
|
0 commit comments