diff --git a/conf/appo/task/go1_joystick_flat/base.yaml b/conf/appo/task/go1_joystick_flat/base.yaml index 230a507e4..8a2f3bf57 100644 --- a/conf/appo/task/go1_joystick_flat/base.yaml +++ b/conf/appo/task/go1_joystick_flat/base.yaml @@ -191,33 +191,24 @@ env: critic_observation_group: critic reward: - # Base-velocity terms bind the IMU-style XML sensors (local_linvel / gyro) - # like the legacy task; Entity root_link_*_vel_b reads base-relative frame - # sensors that are identically zero for the root body on Motrix. tracking_lin_vel: - func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel + func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel_xy_exp weight: 1.0 params: - tracking_sigma: 0.25 + std: 0.5 command_name: twist - sensor_name: local_linvel tracking_ang_vel: - func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel + func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel_z_exp weight: 0.2 params: - tracking_sigma: 0.25 + std: 0.5 command_name: twist - sensor_name: gyro lin_vel_z: - func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z + func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z_l2 weight: -5.0 - params: - sensor_name: local_linvel ang_vel_xy: - func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy + func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy_l2 weight: -0.1 - params: - sensor_name: gyro base_height: func: unilab.tasks.locomotion.common.manager_terms.base_height_l2 weight: -100.0 diff --git a/conf/appo/task/go2_joystick_flat/base.yaml b/conf/appo/task/go2_joystick_flat/base.yaml index 56f4d0fa9..be7364f78 100644 --- a/conf/appo/task/go2_joystick_flat/base.yaml +++ b/conf/appo/task/go2_joystick_flat/base.yaml @@ -154,33 +154,24 @@ env: critic_observation_group: critic reward: - # Base-velocity terms bind the IMU-style XML sensors (local_linvel / gyro) - # like the legacy task; Entity root_link_*_vel_b reads base-relative frame - # sensors that are identically zero for the root body on Motrix. tracking_lin_vel: - func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel + func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel_xy_exp weight: 1.0 params: - tracking_sigma: 0.25 + std: 0.5 command_name: twist - sensor_name: local_linvel tracking_ang_vel: - func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel + func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel_z_exp weight: 0.2 params: - tracking_sigma: 0.25 + std: 0.5 command_name: twist - sensor_name: gyro lin_vel_z: - func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z + func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z_l2 weight: -5.0 - params: - sensor_name: local_linvel ang_vel_xy: - func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy + func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy_l2 weight: -0.1 - params: - sensor_name: gyro base_height: func: unilab.tasks.locomotion.common.manager_terms.base_height_l2 weight: -100.0 diff --git a/conf/flashsac/task/go2_joystick_flat/base.yaml b/conf/flashsac/task/go2_joystick_flat/base.yaml index fb03e18b6..409129c16 100644 --- a/conf/flashsac/task/go2_joystick_flat/base.yaml +++ b/conf/flashsac/task/go2_joystick_flat/base.yaml @@ -154,33 +154,24 @@ env: critic_observation_group: critic reward: - # Base-velocity terms bind the IMU-style XML sensors (local_linvel / gyro) - # like the legacy task; Entity root_link_*_vel_b reads base-relative frame - # sensors that are identically zero for the root body on Motrix. tracking_lin_vel: - func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel + func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel_xy_exp weight: 1.0 params: - tracking_sigma: 0.25 + std: 0.5 command_name: twist - sensor_name: local_linvel tracking_ang_vel: - func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel + func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel_z_exp weight: 0.2 params: - tracking_sigma: 0.25 + std: 0.5 command_name: twist - sensor_name: gyro lin_vel_z: - func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z + func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z_l2 weight: -5.0 - params: - sensor_name: local_linvel ang_vel_xy: - func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy + func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy_l2 weight: -0.1 - params: - sensor_name: gyro base_height: func: unilab.tasks.locomotion.common.manager_terms.base_height_l2 weight: -100.0 diff --git a/conf/flashsac/task/go2_joystick_flat/mujoco.yaml b/conf/flashsac/task/go2_joystick_flat/mujoco.yaml index 47276c2fe..7fbfdec78 100644 --- a/conf/flashsac/task/go2_joystick_flat/mujoco.yaml +++ b/conf/flashsac/task/go2_joystick_flat/mujoco.yaml @@ -97,13 +97,11 @@ reward: tracking_lin_vel: weight: 1.0 params: - # exp(-err / tracking_sigma) with sigma 0.4 matches the former - # exp(-err / std**2) with std 0.6324555320336759. - tracking_sigma: 0.4 + std: 0.6324555320336759 tracking_ang_vel: weight: 0.2 params: - tracking_sigma: 0.4 + std: 0.6324555320336759 lin_vel_z: weight: -5.0 ang_vel_xy: diff --git a/conf/ppo/task/a2_joystick_flat/base.yaml b/conf/ppo/task/a2_joystick_flat/base.yaml index d7f8c8b6c..87203212a 100644 --- a/conf/ppo/task/a2_joystick_flat/base.yaml +++ b/conf/ppo/task/a2_joystick_flat/base.yaml @@ -222,33 +222,24 @@ env: critic_observation_group: critic reward: - # Base-velocity terms bind the IMU-style XML sensors (local_linvel / gyro) - # like the legacy task; Entity root_link_*_vel_b reads base-relative frame - # sensors that are identically zero for the root body on Motrix. tracking_lin_vel: - func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel + func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel_xy_exp weight: 1.0 params: - tracking_sigma: 0.25 + std: 0.5 command_name: twist - sensor_name: local_linvel tracking_ang_vel: - func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel + func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel_z_exp weight: 0.4 params: - tracking_sigma: 0.25 + std: 0.5 command_name: twist - sensor_name: gyro lin_vel_z: - func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z + func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z_l2 weight: -5.0 - params: - sensor_name: local_linvel ang_vel_xy: - func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy + func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy_l2 weight: -0.1 - params: - sensor_name: gyro base_height: func: unilab.tasks.locomotion.common.manager_terms.base_height_l2 weight: -100.0 diff --git a/conf/ppo/task/go1_joystick_flat/base.yaml b/conf/ppo/task/go1_joystick_flat/base.yaml index 230a507e4..8a2f3bf57 100644 --- a/conf/ppo/task/go1_joystick_flat/base.yaml +++ b/conf/ppo/task/go1_joystick_flat/base.yaml @@ -191,33 +191,24 @@ env: critic_observation_group: critic reward: - # Base-velocity terms bind the IMU-style XML sensors (local_linvel / gyro) - # like the legacy task; Entity root_link_*_vel_b reads base-relative frame - # sensors that are identically zero for the root body on Motrix. tracking_lin_vel: - func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel + func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel_xy_exp weight: 1.0 params: - tracking_sigma: 0.25 + std: 0.5 command_name: twist - sensor_name: local_linvel tracking_ang_vel: - func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel + func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel_z_exp weight: 0.2 params: - tracking_sigma: 0.25 + std: 0.5 command_name: twist - sensor_name: gyro lin_vel_z: - func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z + func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z_l2 weight: -5.0 - params: - sensor_name: local_linvel ang_vel_xy: - func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy + func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy_l2 weight: -0.1 - params: - sensor_name: gyro base_height: func: unilab.tasks.locomotion.common.manager_terms.base_height_l2 weight: -100.0 diff --git a/conf/ppo/task/go2_joystick_flat/base.yaml b/conf/ppo/task/go2_joystick_flat/base.yaml index fb03e18b6..409129c16 100644 --- a/conf/ppo/task/go2_joystick_flat/base.yaml +++ b/conf/ppo/task/go2_joystick_flat/base.yaml @@ -154,33 +154,24 @@ env: critic_observation_group: critic reward: - # Base-velocity terms bind the IMU-style XML sensors (local_linvel / gyro) - # like the legacy task; Entity root_link_*_vel_b reads base-relative frame - # sensors that are identically zero for the root body on Motrix. tracking_lin_vel: - func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel + func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel_xy_exp weight: 1.0 params: - tracking_sigma: 0.25 + std: 0.5 command_name: twist - sensor_name: local_linvel tracking_ang_vel: - func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel + func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel_z_exp weight: 0.2 params: - tracking_sigma: 0.25 + std: 0.5 command_name: twist - sensor_name: gyro lin_vel_z: - func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z + func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z_l2 weight: -5.0 - params: - sensor_name: local_linvel ang_vel_xy: - func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy + func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy_l2 weight: -0.1 - params: - sensor_name: gyro base_height: func: unilab.tasks.locomotion.common.manager_terms.base_height_l2 weight: -100.0 diff --git a/conf/ppo/task/go2w_joystick_flat/base.yaml b/conf/ppo/task/go2w_joystick_flat/base.yaml index 5ac4ff4ce..3d057e1ee 100644 --- a/conf/ppo/task/go2w_joystick_flat/base.yaml +++ b/conf/ppo/task/go2w_joystick_flat/base.yaml @@ -203,33 +203,24 @@ env: critic_observation_group: critic reward: - # Base-velocity terms bind the IMU-style XML sensors (local_linvel / gyro) - # like the legacy task; Entity root_link_*_vel_b reads base-relative frame - # sensors that are identically zero for the root body on Motrix. tracking_lin_vel: - func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel + func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel_xy_exp weight: 1.0 params: - tracking_sigma: 0.25 + std: 0.5 command_name: twist - sensor_name: local_linvel tracking_ang_vel: - func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel + func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel_z_exp weight: 0.75 params: - tracking_sigma: 0.25 + std: 0.5 command_name: twist - sensor_name: gyro lin_vel_z: - func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z + func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z_l2 weight: -5.0 - params: - sensor_name: local_linvel ang_vel_xy: - func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy + func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy_l2 weight: -0.1 - params: - sensor_name: gyro base_height: func: unilab.tasks.locomotion.common.manager_terms.base_height_l2 weight: -100.0 diff --git a/conf/ppo/task/quadruped_joystick_rough/base.yaml b/conf/ppo/task/quadruped_joystick_rough/base.yaml index 3c30828cd..c916ccd63 100644 --- a/conf/ppo/task/quadruped_joystick_rough/base.yaml +++ b/conf/ppo/task/quadruped_joystick_rough/base.yaml @@ -129,33 +129,24 @@ env: critic_observation_group: critic reward: - # Base-velocity terms bind the IMU-style XML sensors (local_linvel / gyro) - # like the legacy task; Entity root_link_*_vel_b reads base-relative frame - # sensors that are identically zero for the root body on Motrix. tracking_lin_vel: - func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel + func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel_xy_exp weight: 3.0 params: - tracking_sigma: 0.25 + std: 0.5 command_name: twist - sensor_name: local_linvel tracking_ang_vel: - func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel + func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel_z_exp weight: 1.5 params: - tracking_sigma: 0.25 + std: 0.5 command_name: twist - sensor_name: gyro lin_vel_z: - func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z + func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z_l2 weight: -2.0 - params: - sensor_name: local_linvel ang_vel_xy: - func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy + func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy_l2 weight: -0.05 - params: - sensor_name: gyro action_rate: func: unilab.envs.mdp.action_rate_l2 weight: -0.01 diff --git a/conf/sac/task/go2_joystick_flat/base.yaml b/conf/sac/task/go2_joystick_flat/base.yaml index fb03e18b6..409129c16 100644 --- a/conf/sac/task/go2_joystick_flat/base.yaml +++ b/conf/sac/task/go2_joystick_flat/base.yaml @@ -154,33 +154,24 @@ env: critic_observation_group: critic reward: - # Base-velocity terms bind the IMU-style XML sensors (local_linvel / gyro) - # like the legacy task; Entity root_link_*_vel_b reads base-relative frame - # sensors that are identically zero for the root body on Motrix. tracking_lin_vel: - func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel + func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel_xy_exp weight: 1.0 params: - tracking_sigma: 0.25 + std: 0.5 command_name: twist - sensor_name: local_linvel tracking_ang_vel: - func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel + func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel_z_exp weight: 0.2 params: - tracking_sigma: 0.25 + std: 0.5 command_name: twist - sensor_name: gyro lin_vel_z: - func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z + func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z_l2 weight: -5.0 - params: - sensor_name: local_linvel ang_vel_xy: - func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy + func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy_l2 weight: -0.1 - params: - sensor_name: gyro base_height: func: unilab.tasks.locomotion.common.manager_terms.base_height_l2 weight: -100.0 diff --git a/conf/sac/task/go2w_joystick_flat/base.yaml b/conf/sac/task/go2w_joystick_flat/base.yaml index 5ac4ff4ce..3d057e1ee 100644 --- a/conf/sac/task/go2w_joystick_flat/base.yaml +++ b/conf/sac/task/go2w_joystick_flat/base.yaml @@ -203,33 +203,24 @@ env: critic_observation_group: critic reward: - # Base-velocity terms bind the IMU-style XML sensors (local_linvel / gyro) - # like the legacy task; Entity root_link_*_vel_b reads base-relative frame - # sensors that are identically zero for the root body on Motrix. tracking_lin_vel: - func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel + func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel_xy_exp weight: 1.0 params: - tracking_sigma: 0.25 + std: 0.5 command_name: twist - sensor_name: local_linvel tracking_ang_vel: - func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel + func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel_z_exp weight: 0.75 params: - tracking_sigma: 0.25 + std: 0.5 command_name: twist - sensor_name: gyro lin_vel_z: - func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z + func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z_l2 weight: -5.0 - params: - sensor_name: local_linvel ang_vel_xy: - func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy + func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy_l2 weight: -0.1 - params: - sensor_name: gyro base_height: func: unilab.tasks.locomotion.common.manager_terms.base_height_l2 weight: -100.0 diff --git a/conf/td3/task/go1_joystick_flat/base.yaml b/conf/td3/task/go1_joystick_flat/base.yaml index 230a507e4..8a2f3bf57 100644 --- a/conf/td3/task/go1_joystick_flat/base.yaml +++ b/conf/td3/task/go1_joystick_flat/base.yaml @@ -191,33 +191,24 @@ env: critic_observation_group: critic reward: - # Base-velocity terms bind the IMU-style XML sensors (local_linvel / gyro) - # like the legacy task; Entity root_link_*_vel_b reads base-relative frame - # sensors that are identically zero for the root body on Motrix. tracking_lin_vel: - func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel + func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel_xy_exp weight: 1.0 params: - tracking_sigma: 0.25 + std: 0.5 command_name: twist - sensor_name: local_linvel tracking_ang_vel: - func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel + func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel_z_exp weight: 0.2 params: - tracking_sigma: 0.25 + std: 0.5 command_name: twist - sensor_name: gyro lin_vel_z: - func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z + func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z_l2 weight: -5.0 - params: - sensor_name: local_linvel ang_vel_xy: - func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy + func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy_l2 weight: -0.1 - params: - sensor_name: gyro base_height: func: unilab.tasks.locomotion.common.manager_terms.base_height_l2 weight: -100.0 diff --git a/conf/td3/task/go2_joystick_flat/base.yaml b/conf/td3/task/go2_joystick_flat/base.yaml index fb03e18b6..409129c16 100644 --- a/conf/td3/task/go2_joystick_flat/base.yaml +++ b/conf/td3/task/go2_joystick_flat/base.yaml @@ -154,33 +154,24 @@ env: critic_observation_group: critic reward: - # Base-velocity terms bind the IMU-style XML sensors (local_linvel / gyro) - # like the legacy task; Entity root_link_*_vel_b reads base-relative frame - # sensors that are identically zero for the root body on Motrix. tracking_lin_vel: - func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel + func: unilab.tasks.locomotion.common.manager_terms.track_lin_vel_xy_exp weight: 1.0 params: - tracking_sigma: 0.25 + std: 0.5 command_name: twist - sensor_name: local_linvel tracking_ang_vel: - func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel + func: unilab.tasks.locomotion.common.manager_terms.track_ang_vel_z_exp weight: 0.2 params: - tracking_sigma: 0.25 + std: 0.5 command_name: twist - sensor_name: gyro lin_vel_z: - func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z + func: unilab.tasks.locomotion.common.manager_terms.lin_vel_z_l2 weight: -5.0 - params: - sensor_name: local_linvel ang_vel_xy: - func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy + func: unilab.tasks.locomotion.common.manager_terms.ang_vel_xy_l2 weight: -0.1 - params: - sensor_name: gyro base_height: func: unilab.tasks.locomotion.common.manager_terms.base_height_l2 weight: -100.0 diff --git a/src/unilab/base/backend/base.py b/src/unilab/base/backend/base.py index 211811b76..081da0f5d 100644 --- a/src/unilab/base/backend/base.py +++ b/src/unilab/base/backend/base.py @@ -851,7 +851,14 @@ def get_body_quat_b(self, body_ids: np.ndarray) -> np.ndarray: @abc.abstractmethod def get_body_lin_vel_b(self, body_ids: np.ndarray) -> np.ndarray: - """Return selected body linear velocities in the baselink frame. + """Return selected body linear velocities expressed in each body's own frame. + + The value is the body's world-frame velocity rotated by the inverse of + the body's world-frame orientation, i.e. + ``quat_apply_inverse(quat_w, lin_vel_w)`` (mjlab/Isaac-style analytical + definition). It is well-defined for every body — including the root + body — and must NOT be implemented as the motion relative to the + baselink frame (which degenerates to zero for the root body). Args: body_ids: Body ID array. @@ -862,7 +869,14 @@ def get_body_lin_vel_b(self, body_ids: np.ndarray) -> np.ndarray: @abc.abstractmethod def get_body_ang_vel_b(self, body_ids: np.ndarray) -> np.ndarray: - """Return selected body angular velocities in the baselink frame. + """Return selected body angular velocities expressed in each body's own frame. + + The value is the body's world-frame angular velocity rotated by the + inverse of the body's world-frame orientation, i.e. + ``quat_apply_inverse(quat_w, ang_vel_w)`` (mjlab/Isaac-style analytical + definition). It is well-defined for every body — including the root + body — and must NOT be implemented as the motion relative to the + baselink frame (which degenerates to zero for the root body). Args: body_ids: Body ID array. diff --git a/src/unilab/base/backend/drake/backend.py b/src/unilab/base/backend/drake/backend.py index e58251c44..b508950af 100644 --- a/src/unilab/base/backend/drake/backend.py +++ b/src/unilab/base/backend/drake/backend.py @@ -713,19 +713,23 @@ def get_body_quat_b(self, body_ids: np.ndarray) -> np.ndarray: return _quat_multiply(base_inv[:, None, :], body_quat) def get_body_lin_vel_b(self, body_ids: np.ndarray) -> np.ndarray: - base_rot = _quat_to_rotation_matrix(self.get_base_quat()) + # Analytical per the SimBackend contract: world-frame velocity + # expressed in each body's own frame. + body_state = self._body_state(body_ids) + body_rot = _quat_to_rotation_matrix(body_state["quat"]) return np.einsum( - "nij,nkj->nki", - np.swapaxes(base_rot, 1, 2), - self._body_state(body_ids)["linvel"], + "nkij,nkj->nki", + np.swapaxes(body_rot, -1, -2), + body_state["linvel"], ) def get_body_ang_vel_b(self, body_ids: np.ndarray) -> np.ndarray: - base_rot = _quat_to_rotation_matrix(self.get_base_quat()) + body_state = self._body_state(body_ids) + body_rot = _quat_to_rotation_matrix(body_state["quat"]) return np.einsum( - "nij,nkj->nki", - np.swapaxes(base_rot, 1, 2), - self._body_state(body_ids)["angvel"], + "nkij,nkj->nki", + np.swapaxes(body_rot, -1, -2), + body_state["angvel"], ) def get_sensor_data(self, name: str) -> np.ndarray: diff --git a/src/unilab/base/backend/motrix/backend.py b/src/unilab/base/backend/motrix/backend.py index 53c02ec97..53418dfe6 100644 --- a/src/unilab/base/backend/motrix/backend.py +++ b/src/unilab/base/backend/motrix/backend.py @@ -22,6 +22,7 @@ IntervalRandomizationPlan, ResetRandomizationPayload, ) +from unilab.utils.rotation import np_quat_apply_inverse_batched try: import motrixsim as mtx @@ -1175,10 +1176,19 @@ def get_body_quat_b(self, body_ids: np.ndarray) -> np.ndarray: return self._xyzw_to_wxyz(self._get_body_sensor_values(body_ids, "track_quat_b")) def get_body_lin_vel_b(self, body_ids: np.ndarray) -> np.ndarray: - return self._get_body_sensor_values(body_ids, "track_linvel_b") + # Analytical per the SimBackend contract: world-frame velocity rotated + # into each body's own frame. MotrixSim frame sensors report motion + # relative to the baselink and degenerate to zero for the root body. + ids = self._as_body_ids(body_ids) + return np_quat_apply_inverse_batched( + self.get_body_quat_w(ids), self._get_link_lin_vel_w(ids) + ) def get_body_ang_vel_b(self, body_ids: np.ndarray) -> np.ndarray: - return self._get_body_sensor_values(body_ids, "track_angvel_b") + ids = self._as_body_ids(body_ids) + return np_quat_apply_inverse_batched( + self.get_body_quat_w(ids), self._get_link_ang_vel_w(ids) + ) # ------------------------------------------------------------------ # # Sensors # diff --git a/src/unilab/base/backend/motrix/scene.py b/src/unilab/base/backend/motrix/scene.py index 2d65d1208..b4a8bc432 100644 --- a/src/unilab/base/backend/motrix/scene.py +++ b/src/unilab/base/backend/motrix/scene.py @@ -132,7 +132,14 @@ def _motrix_world_link_names(world: World) -> list[str]: def add_motrix_tracking_frame_sensors(world: World, *, base_name: str) -> None: - """Add Motrix-native frame sensors matching the legacy tracking sensor contract.""" + """Add Motrix-native frame sensors matching the legacy tracking sensor contract. + + Only pose sensors are added: body-frame velocities are computed + analytically from the world-frame link state (see + ``MotrixBackend.get_body_*_vel_b``), because MotrixSim frame velocity + sensors report motion relative to the baselink and degenerate to zero for + the root body. + """ import motrixsim.msd as msd link_names = _motrix_world_link_names(world) @@ -143,8 +150,6 @@ def add_motrix_tracking_frame_sensors(world: World, *, base_name: str) -> None: sensor_specs = ( ("track_pos_b", msd.FrameSensorType.FramePos), ("track_quat_b", msd.FrameSensorType.FrameQuat), - ("track_linvel_b", msd.FrameSensorType.FrameLinVel), - ("track_angvel_b", msd.FrameSensorType.FrameAngVel), ) ref_frame = msd.FrameSensorRef.object(msd.ObjectType.link(base_name)) for link_name in link_names: diff --git a/src/unilab/base/backend/mujoco/backend.py b/src/unilab/base/backend/mujoco/backend.py index 6069a6fe0..5aa13fd07 100644 --- a/src/unilab/base/backend/mujoco/backend.py +++ b/src/unilab/base/backend/mujoco/backend.py @@ -32,6 +32,7 @@ ResetRandomizationPayload, ) from unilab.dtype_config import get_global_dtype +from unilab.utils.rotation import np_quat_apply_inverse_batched from ..base import ( BackendHeightScanner, @@ -437,8 +438,6 @@ def _get_sensor_view(prefix, dim): # Local (baselink) sensors self._tracked_pos_b_all = _get_sensor_view("track_pos_b", 3) self._tracked_quat_b_all = _get_sensor_view("track_quat_b", 4) - self._tracked_linvel_b_all = _get_sensor_view("track_linvel_b", 3) - self._tracked_angvel_b_all = _get_sensor_view("track_angvel_b", 3) def _load_base_model(self) -> mujoco.MjModel: if isinstance(self._model_file, mujoco.MjModel): @@ -1438,10 +1437,19 @@ def get_body_quat_b(self, body_ids: np.ndarray) -> np.ndarray: return self._tracked_quat_b_all[:, self._get_mapped_indices(body_ids), :] # type: ignore[no-any-return] def get_body_lin_vel_b(self, body_ids: np.ndarray) -> np.ndarray: - return self._tracked_linvel_b_all[:, self._get_mapped_indices(body_ids), :] # type: ignore[no-any-return] + # Analytical per the SimBackend contract: world-frame velocity rotated + # into each body's own frame. MuJoCo framelinvel sensors with a baselink + # reference report relative motion and degenerate to zero for the root. + idx = self._get_mapped_indices(body_ids) + return np_quat_apply_inverse_batched( + self._tracked_quat_w_all[:, idx, :], self._tracked_linvel_w_all[:, idx, :] + ) def get_body_ang_vel_b(self, body_ids: np.ndarray) -> np.ndarray: - return self._tracked_angvel_b_all[:, self._get_mapped_indices(body_ids), :] # type: ignore[no-any-return] + idx = self._get_mapped_indices(body_ids) + return np_quat_apply_inverse_batched( + self._tracked_quat_w_all[:, idx, :], self._tracked_angvel_w_all[:, idx, :] + ) # ------------------------------------------------------------------ # # Sensors # diff --git a/src/unilab/base/backend/mujoco/xml.py b/src/unilab/base/backend/mujoco/xml.py index 62b1d2555..48a5d7897 100644 --- a/src/unilab/base/backend/mujoco/xml.py +++ b/src/unilab/base/backend/mujoco/xml.py @@ -133,6 +133,10 @@ def _add_w_sensors(spec, valid_bnames: list[str]) -> None: def _add_b_sensors(spec, valid_bnames: list[str], baselink_name: str) -> None: + # Body-frame velocities are computed analytically from the world-frame + # tracking sensors (see MuJoCoBackend.get_body_*_vel_b); framelinvel / + # frameangvel sensors with a baselink reference report relative motion and + # degenerate to zero for the root body, so they are intentionally absent. mujoco = _mujoco_module() for bname in valid_bnames: spec.add_sensor( @@ -152,24 +156,6 @@ def _add_b_sensors(spec, valid_bnames: list[str], baselink_name: str) -> None: reftype=mujoco.mjtObj.mjOBJ_XBODY, refname=baselink_name, ) - for bname in valid_bnames: - spec.add_sensor( - name=f"track_linvel_b_{bname}", - type=mujoco.mjtSensor.mjSENS_FRAMELINVEL, - objtype=mujoco.mjtObj.mjOBJ_XBODY, - objname=bname, - reftype=mujoco.mjtObj.mjOBJ_XBODY, - refname=baselink_name, - ) - for bname in valid_bnames: - spec.add_sensor( - name=f"track_angvel_b_{bname}", - type=mujoco.mjtSensor.mjSENS_FRAMEANGVEL, - objtype=mujoco.mjtObj.mjOBJ_XBODY, - objname=bname, - reftype=mujoco.mjtObj.mjOBJ_XBODY, - refname=baselink_name, - ) def _write_temp_xml(tree: ET.ElementTree[ET.Element], model_file: str) -> str: # type: ignore[type-arg] diff --git a/src/unilab/tasks/locomotion/common/manager_terms.py b/src/unilab/tasks/locomotion/common/manager_terms.py index c7b8212d9..25e743c1a 100644 --- a/src/unilab/tasks/locomotion/common/manager_terms.py +++ b/src/unilab/tasks/locomotion/common/manager_terms.py @@ -13,7 +13,6 @@ from unilab.dtype_config import get_global_dtype from unilab.managers.manager_base import ManagerTermBase, ManagerTermBaseCfg from unilab.managers.scene_entity_config import SceneEntityCfg -from unilab.tasks.locomotion.common.sensor_terms import SensorTermBase if TYPE_CHECKING: from unilab.base.entity import Entity @@ -231,116 +230,6 @@ def stand_still_l1( ) -# --------------------------------------------------------------------------- -# Sensor-bound base velocity terms (legacy quadruped equations) -# --------------------------------------------------------------------------- -# -# These terms read the named IMU-style XML sensors (``local_linvel`` / -# ``gyro``) exactly like the legacy Go1/Go2 joystick tasks. They deliberately -# avoid ``Entity.data.root_link_*_vel_b``: on backends that realize body-frame -# velocity through base-relative frame sensors (Motrix), the root body's value -# is identically zero, which silently zeroes every base-velocity reward term. - - -class _NamedSensorTerm(SensorTermBase): - """Single-sensor binding with a configurable ``sensor_name`` parameter.""" - - _allowed_params: ClassVar[frozenset[str]] = frozenset({"sensor_name"}) - _default_sensor_name: ClassVar[str] = "" - - def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): - super().__init__(cfg, env) - sensor_name = cfg.params.get("sensor_name", self._default_sensor_name) - if not isinstance(sensor_name, str) or not sensor_name: - raise ValueError(f"{self.name} sensor_name must be a non-empty string") - self._sensor_name = sensor_name - self._view = self._bind((sensor_name,)) - - def _read_sensor(self, env: ManagerBasedRlEnv) -> np.ndarray: - value = self._read(self._view, self.name) - return _state(self.name, f"sensor '{self._sensor_name}'", value, (env.num_envs, 3)) - - -class _LinVelSensorTerm(_NamedSensorTerm): - _default_sensor_name: ClassVar[str] = "local_linvel" - - -class _GyroSensorTerm(_NamedSensorTerm): - _default_sensor_name: ClassVar[str] = "gyro" - - -class track_lin_vel(_LinVelSensorTerm): - """Exponential reward for tracking commanded xy linear velocity.""" - - _allowed_params = _NamedSensorTerm._allowed_params | {"tracking_sigma", "command_name"} - - def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): - super().__init__(cfg, env) - self._sigma = _real( - self.name, - "tracking_sigma", - cfg.params.get("tracking_sigma", 0.25), - minimum=0.0, - strict_minimum=True, - ) - command_name = cfg.params.get("command_name", "twist") - if not isinstance(command_name, str) or not command_name: - raise ValueError(f"{self.name} command_name must be a non-empty string") - self._command_name = command_name - - def __call__(self, env: ManagerBasedRlEnv, **params: Any) -> np.ndarray: - del params - linvel = self._read_sensor(env) - command = _command(env, self.name, self._command_name) - error = np.sum(np.square(command[:, :2] - linvel[:, :2]), axis=1) - return np.asarray(np.exp(-error / self._sigma), dtype=get_global_dtype()) - - -class track_ang_vel(_GyroSensorTerm): - """Exponential reward for tracking commanded yaw angular velocity.""" - - _allowed_params = _NamedSensorTerm._allowed_params | {"tracking_sigma", "command_name"} - - def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): - super().__init__(cfg, env) - self._sigma = _real( - self.name, - "tracking_sigma", - cfg.params.get("tracking_sigma", 0.25), - minimum=0.0, - strict_minimum=True, - ) - command_name = cfg.params.get("command_name", "twist") - if not isinstance(command_name, str) or not command_name: - raise ValueError(f"{self.name} command_name must be a non-empty string") - self._command_name = command_name - - def __call__(self, env: ManagerBasedRlEnv, **params: Any) -> np.ndarray: - del params - gyro = self._read_sensor(env) - command = _command(env, self.name, self._command_name) - error = np.square(command[:, 2] - gyro[:, 2]) - return np.asarray(np.exp(-error / self._sigma), dtype=get_global_dtype()) - - -class lin_vel_z(_LinVelSensorTerm): - """Penalty for vertical (z) linear velocity.""" - - def __call__(self, env: ManagerBasedRlEnv, **params: Any) -> np.ndarray: - del params - linvel = self._read_sensor(env) - return np.asarray(np.square(linvel[:, 2]), dtype=get_global_dtype()) - - -class ang_vel_xy(_GyroSensorTerm): - """Penalty for roll/pitch angular velocity.""" - - def __call__(self, env: ManagerBasedRlEnv, **params: Any) -> np.ndarray: - del params - gyro = self._read_sensor(env) - return np.asarray(np.sum(np.square(gyro[:, :2]), axis=1), dtype=get_global_dtype()) - - class _GaitTerm(ManagerTermBase): _allowed_params: ClassVar[frozenset[str]] = frozenset( {"frequency", "phase_offsets", "command_name", "command_threshold"} @@ -584,19 +473,15 @@ def __call__(self, env: ManagerBasedRlEnv, **params: Any) -> np.ndarray: __all__ = [ - "ang_vel_xy", "ang_vel_xy_l2", "base_height_l2", "feet_air_while_standing", "feet_phase_contact", "feet_phase_swing_height", "joint_deviation_l1", - "lin_vel_z", "lin_vel_z_l2", "quadruped_gait_phase", "stand_still_l1", - "track_ang_vel", "track_ang_vel_z_exp", - "track_lin_vel", "track_lin_vel_xy_exp", ] diff --git a/src/unilab/utils/rotation.py b/src/unilab/utils/rotation.py index 2879a02fb..aaa39dc2f 100644 --- a/src/unilab/utils/rotation.py +++ b/src/unilab/utils/rotation.py @@ -64,11 +64,11 @@ def np_quat_mul_batched(q1: np.ndarray, q2: np.ndarray) -> np.ndarray: def np_quat_conjugate(q: np.ndarray) -> np.ndarray: - """Conjugate of unit quaternions (N, 4) or (4,), w-first.""" + """Conjugate of unit quaternions (..., 4), w-first.""" if q.ndim == 1: return np.array([q[0], -q[1], -q[2], -q[3]]) conj = q.copy() - conj[:, 1:] *= -1 + conj[..., 1:] *= -1 return conj # type: ignore[no-any-return] @@ -258,6 +258,15 @@ def np_quat_apply_inverse(q: np.ndarray, v: np.ndarray) -> np.ndarray: return np_quat_apply(np_quat_inv(q), v) +def np_quat_apply_inverse_batched(q: np.ndarray, v: np.ndarray) -> np.ndarray: + """Rotate broadcast-compatible vector arrays by inverse quaternions. + + ``q`` has shape (..., 4), ``v`` has shape (..., 3), and leading dimensions + are broadcast. Batched counterpart of ``np_quat_apply_inverse``. + """ + return np_quat_apply_batched(np_quat_conjugate(q), v) + + def np_quat_error_magnitude(q1: np.ndarray, q2: np.ndarray) -> np.ndarray: """Angular error magnitude between quaternions (N,) or scalar.""" q1_was_1d = q1.ndim == 1 diff --git a/tests/base/test_sim_backend.py b/tests/base/test_sim_backend.py index 17996fe11..0647c29e3 100644 --- a/tests/base/test_sim_backend.py +++ b/tests/base/test_sim_backend.py @@ -93,6 +93,52 @@ def _allegro_state() -> tuple[np.ndarray, np.ndarray]: return qpos, qvel +# Base-motion state used by the get_body_*_vel_b semantic tests: a 30-degree +# yawed base with a non-zero velocity. Free-joint convention (shared by the +# backends' set_state): qvel[:3] is linear velocity in the world frame, +# qvel[3:6] is angular velocity in the body-local frame. +_BASE_YAW = np.deg2rad(30.0) +_BASE_LIN_VEL_W = np.array([0.5, -0.2, 0.8]) +_BASE_ANG_VEL_B = np.array([0.3, -0.1, 0.2]) + + +def _tilted_motion_state(bkd) -> tuple[np.ndarray, np.ndarray]: + nq = bkd.get_dof_pos().shape[-1] + 7 + nv = bkd.get_dof_vel().shape[-1] + 6 + qpos = np.tile(_identity_qpos_mujoco(nq), (NUM_ENVS, 1)) + qpos[:, 3] = np.cos(_BASE_YAW / 2) + qpos[:, 6] = np.sin(_BASE_YAW / 2) + qvel = np.zeros((NUM_ENVS, nv)) + qvel[:, :3] = _BASE_LIN_VEL_W + qvel[:, 3:6] = _BASE_ANG_VEL_B + return qpos, qvel + + +def _yaw_world_to_body_rot() -> np.ndarray: + c, s = np.cos(_BASE_YAW), np.sin(_BASE_YAW) + return np.array([[c, s, 0.0], [-s, c, 0.0], [0.0, 0.0, 1.0]]) + + +def _assert_root_vel_b_matches_world_vel_in_body_frame(bkd, root_id: int, atol: float) -> None: + """Root-body vel_b must be the world-frame velocity expressed in the body + frame (mjlab/Isaac-style analytical contract), not degenerate zero.""" + qpos, qvel = _tilted_motion_state(bkd) + bkd.set_state(np.arange(NUM_ENVS), qpos, qvel) + ids = np.array([root_id]) + rot_w2b = _yaw_world_to_body_rot() + np.testing.assert_allclose( + bkd.get_body_lin_vel_b(ids)[:, 0, :], + np.tile(rot_w2b @ _BASE_LIN_VEL_W, (NUM_ENVS, 1)), + atol=atol, + ) + # qvel[3:6] is already body-local, so ang_vel_b must reproduce it exactly. + np.testing.assert_allclose( + bkd.get_body_ang_vel_b(ids)[:, 0, :], + np.tile(_BASE_ANG_VEL_B, (NUM_ENVS, 1)), + atol=atol, + ) + + # --------------------------------------------------------------------------- # MuJoCo — basic, 3 robots, no body sensors # --------------------------------------------------------------------------- @@ -583,6 +629,12 @@ def test_get_body_lin_vel_b_shape(self, bkd, body_ids): def test_get_body_ang_vel_b_shape(self, bkd, body_ids): _shape(bkd.get_body_ang_vel_b(body_ids), NUM_ENVS, len(body_ids), 3) + def test_root_body_vel_b_matches_world_vel_in_body_frame(self, bkd): + mujoco = _mujoco_module() + + root_id = int(mujoco.mj_name2id(bkd.model, mujoco.mjtObj.mjOBJ_BODY, _G1["base_name"])) + _assert_root_vel_b_matches_world_vel_in_body_frame(bkd, root_id, atol=1e-5) + # sensors def test_get_sensor_data_w_shape(self, bkd, body_ids): @@ -887,6 +939,10 @@ def test_get_body_lin_vel_b_shape(self, bkd, body_ids): def test_get_body_ang_vel_b_shape(self, bkd, body_ids): _shape(bkd.get_body_ang_vel_b(body_ids), NUM_ENVS, len(body_ids), 3) + def test_root_body_vel_b_matches_world_vel_in_body_frame(self, bkd): + root_id = int(bkd.model.get_body_index(_G1["base_name"])) + _assert_root_vel_b_matches_world_vel_in_body_frame(bkd, root_id, atol=1e-4) + # sensors def test_get_sensor_data_b_shape(self, bkd, body_ids): @@ -1094,6 +1150,25 @@ def test_body_ang_vel_b(self, synced, body_pairs): atol=self.ATOL, ) + def test_body_vel_b_parity_under_base_motion(self, synced, body_pairs): + """vel_b parity must hold with non-zero base motion (root included).""" + mj, mx = synced + qpos, qvel = _tilted_motion_state(mj) + env_idx = np.arange(NUM_ENVS) + mj.set_state(env_idx, qpos, qvel) + mx.set_state(env_idx, qpos, qvel) + mj_ids, mx_ids = body_pairs + np.testing.assert_allclose( + mj.get_body_lin_vel_b(mj_ids), + mx.get_body_lin_vel_b(mx_ids), + atol=self.ATOL, + ) + np.testing.assert_allclose( + mj.get_body_ang_vel_b(mj_ids), + mx.get_body_ang_vel_b(mx_ids), + atol=self.ATOL, + ) + # --------------------------------------------------------------------------- # Unified model properties — MuJoCo diff --git a/tests/config/test_reward_injection.py b/tests/config/test_reward_injection.py index 3c3eb3500..96e5df93e 100644 --- a/tests/config/test_reward_injection.py +++ b/tests/config/test_reward_injection.py @@ -40,7 +40,7 @@ def test_resolve_reward_dict_reads_task_reward(): assert reward_dict["tracking_lin_vel"]["weight"] == 1.0 assert reward_dict["tracking_ang_vel"]["weight"] == 0.2 - assert reward_dict["tracking_lin_vel"]["func"].endswith("manager_terms.track_lin_vel") + assert reward_dict["tracking_lin_vel"]["func"].endswith("track_lin_vel_xy_exp") def test_reward_config_conversion(): diff --git a/tests/envs/locomotion/go2/test_manager_based_cfg.py b/tests/envs/locomotion/go2/test_manager_based_cfg.py index 589c3925a..ef52d4bd1 100644 --- a/tests/envs/locomotion/go2/test_manager_based_cfg.py +++ b/tests/envs/locomotion/go2/test_manager_based_cfg.py @@ -17,7 +17,6 @@ from unilab.base.config_adapter import BackendAdapter from unilab.base.config_materialization import apply_cfg_overrides from unilab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg, mdp -from unilab.tasks.locomotion.common import manager_terms ROOT_DIR = Path(__file__).parents[4] CONF_DIR = ROOT_DIR / "conf" @@ -238,18 +237,6 @@ def test_go2_flat_owner_materializes_complete_plain_manager_cfg( if alive_declared: expected_weights["alive"] = 0.0 assert {name: term.weight for name, term in env_cfg.rewards.items()} == expected_weights - # Base-velocity rewards stay bound to the legacy IMU-style XML sensors; - # Entity root_link_*_vel_b is identically zero for the root on Motrix. - assert env_cfg.rewards["tracking_lin_vel"].func is manager_terms.track_lin_vel - assert env_cfg.rewards["tracking_ang_vel"].func is manager_terms.track_ang_vel - assert env_cfg.rewards["lin_vel_z"].func is manager_terms.lin_vel_z - assert env_cfg.rewards["ang_vel_xy"].func is manager_terms.ang_vel_xy - tracking_params = env_cfg.rewards["tracking_lin_vel"].params - assert tracking_params["tracking_sigma"] == pytest.approx(0.25) - assert tracking_params["sensor_name"] == "local_linvel" - assert env_cfg.rewards["tracking_ang_vel"].params["sensor_name"] == "gyro" - assert env_cfg.rewards["lin_vel_z"].params["sensor_name"] == "local_linvel" - assert env_cfg.rewards["ang_vel_xy"].params["sensor_name"] == "gyro" if fixed_command: ranges = env_cfg.commands["twist"].ranges @@ -375,7 +362,7 @@ def test_go2_flat_flashsac_uses_canonical_manager_events_and_numpy_noise() -> No assert group.terms["joint_vel"].noise.n_min == pytest.approx(-0.1) assert group.terms["joint_vel"].noise.n_max == pytest.approx(0.1) - assert env_cfg.rewards["tracking_lin_vel"].params["tracking_sigma"] == pytest.approx(0.4) + assert env_cfg.rewards["tracking_lin_vel"].params["std"] == pytest.approx(0.4**0.5) assert env_cfg.rewards["base_height"].weight == pytest.approx(-20.0) assert env_cfg.rewards["contact"].weight == pytest.approx(1.5) diff --git a/tests/envs/locomotion/test_manager_gait_terms.py b/tests/envs/locomotion/test_manager_gait_terms.py index b30f799ff..d91bf2eb6 100644 --- a/tests/envs/locomotion/test_manager_gait_terms.py +++ b/tests/envs/locomotion/test_manager_gait_terms.py @@ -38,8 +38,6 @@ def __init__(self) -> None: "fr_pos": np.array([[0, 0, 0.2], [0, 0, 0.1]], dtype=np.float32), "rl_pos": np.array([[0, 0, 0.1], [0, 0, 0.2]], dtype=np.float32), "rr_pos": np.array([[0, 0, 0], [0, 0, 0.1]], dtype=np.float32), - "local_linvel": np.array([[0.2, 0.1, -0.3], [0.3, -0.2, 0.5]], dtype=np.float32), - "gyro": np.array([[0.1, -0.2, 0.4], [-0.3, 0.2, -0.1]], dtype=np.float32), } def bind_sensor_data(self, names) -> BackendSensorView: @@ -366,55 +364,3 @@ def test_base_reward_terms_fail_closed_at_nearest_boundary() -> None: manager_terms.joint_deviation_l1, asset_cfg=SceneEntityCfg("missing"), ) - - -def test_sensor_bound_velocity_terms_match_legacy_go2_equations() -> None: - env = _env() - linvel = cast(Any, env).scene.values["local_linvel"] - gyro = cast(Any, env).scene.values["gyro"] - command = cast(Any, env).command_manager.command - sigma = 0.25 - actual = { - "track_lin": _reward_value( - env, manager_terms.track_lin_vel, tracking_sigma=sigma, command_name="twist" - ), - "track_ang": _reward_value( - env, manager_terms.track_ang_vel, tracking_sigma=sigma, command_name="twist" - ), - "lin_z": _reward_value(env, manager_terms.lin_vel_z), - "ang_xy": _reward_value(env, manager_terms.ang_vel_xy), - } - expected = { - "track_lin": np.exp(-np.sum(np.square(command[:, :2] - linvel[:, :2]), axis=1) / sigma), - "track_ang": np.exp(-np.square(command[:, 2] - gyro[:, 2]) / sigma), - "lin_z": np.square(linvel[:, 2]), - "ang_xy": np.sum(np.square(gyro[:, :2]), axis=1), - } - for name in expected: - assert actual[name].shape == (2,) - assert actual[name].dtype == np.dtype(get_global_dtype()) - np.testing.assert_allclose(actual[name], expected[name], rtol=1e-6, atol=1e-7) - - -def test_sensor_bound_velocity_terms_fail_closed_at_nearest_boundary() -> None: - env = _env() - with pytest.raises(KeyError, match="lin_vel_z.*could not be materialized.*missing"): - _reward_value(env, manager_terms.lin_vel_z, sensor_name="missing") - with pytest.raises(ValueError, match="track_lin_vel tracking_sigma must be greater than 0.0"): - _reward_value(env, manager_terms.track_lin_vel, tracking_sigma=0.0) - with pytest.raises(ValueError, match="track_ang_vel sensor_name must be a non-empty string"): - _reward_value(env, manager_terms.track_ang_vel, sensor_name="") - with pytest.raises(KeyError, match="track_lin_vel command capability 'missing'"): - _reward_value(env, manager_terms.track_lin_vel, command_name="missing") - with pytest.raises(TypeError, match="ang_vel_xy received unsupported parameters"): - _reward_value(env, manager_terms.ang_vel_xy, std=0.5) - - scene = _Scene() - manager = RewardManager( - {"parity": RewardTermCfg(func=manager_terms.ang_vel_xy, weight=1.0)}, - _env(scene=scene), - scale_by_dt=False, - ) - scene.values["gyro"][1, 0] = np.nan - with pytest.raises(ValueError, match="ang_vel_xy.*backend 'fake'.*NaN or Inf"): - manager.compute(dt=0.02) diff --git a/tests/utils/test_math_utils.py b/tests/utils/test_math_utils.py index da9c4effa..f596195fb 100644 --- a/tests/utils/test_math_utils.py +++ b/tests/utils/test_math_utils.py @@ -10,6 +10,7 @@ np_quat_angular_velocity, np_quat_apply, np_quat_apply_batched, + np_quat_apply_inverse_batched, np_quat_ensure_continuity, np_quat_error_magnitude, np_quat_error_magnitude_batched, @@ -159,6 +160,34 @@ def test_batched_quaternion_helpers_match_flattened_helpers() -> None: ) +def test_quat_apply_inverse_batched_matches_matrix_transpose() -> None: + """Batched inverse rotation must equal R(q)^T @ v on (env, body) inputs.""" + num_envs = 3 + num_bodies = 4 + + quat = np_quat_from_euler_xyz( + np.linspace(-0.3, 0.4, num_envs * num_bodies), + np.linspace(0.2, -0.15, num_envs * num_bodies), + np.linspace(-0.5, 0.25, num_envs * num_bodies), + ).reshape(num_envs, num_bodies, 4) + vectors = np.linspace(-0.6, 0.7, num_envs * num_bodies * 3).reshape(num_envs, num_bodies, 3) + + mats = np_matrix_from_quat(quat.reshape(-1, 4)) + expected = np.einsum("nji,nj->ni", mats, vectors.reshape(-1, 3)) + + np.testing.assert_allclose( + np_quat_apply_inverse_batched(quat, vectors).reshape(-1, 3), + expected, + atol=1e-12, + ) + # Round-trip: applying q after q^-1 must recover the original vector. + np.testing.assert_allclose( + np_quat_apply_batched(quat, np_quat_apply_inverse_batched(quat, vectors)), + vectors, + atol=1e-12, + ) + + def test_anchor_frame_transform_matches_flattened_path() -> None: num_envs = 3 num_bodies = 4 diff --git a/tests/utils/test_xml_utils.py b/tests/utils/test_xml_utils.py index 27eb03d9b..ee17db92d 100644 --- a/tests/utils/test_xml_utils.py +++ b/tests/utils/test_xml_utils.py @@ -105,8 +105,6 @@ def test_inject_mujoco_tracking_sensors_uses_mjspec_and_preserves_contract() -> "track_angvel_w_pelvis", "track_pos_b_pelvis", "track_quat_b_pelvis", - "track_linvel_b_pelvis", - "track_angvel_b_pelvis", ): assert mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, sensor_name) >= 0 finally: @@ -126,8 +124,6 @@ def test_materialize_motrix_scene_adds_tracking_frame_sensors() -> None: for sensor_name in ( "track_pos_b_pelvis", "track_quat_b_pelvis", - "track_linvel_b_pelvis", - "track_angvel_b_pelvis", ): assert model.get_sensor_value(sensor_name, data).shape[0] == 1