diff --git a/conf/ppo/task/g1_motion_tracking/motrix.yaml b/conf/ppo/task/g1_motion_tracking/motrix.yaml index 2a3343635..1b742d21f 100644 --- a/conf/ppo/task/g1_motion_tracking/motrix.yaml +++ b/conf/ppo/task/g1_motion_tracking/motrix.yaml @@ -38,6 +38,7 @@ reward: motion_joint_vel: 0.0 action_rate_l2: -0.05 joint_limit: -10.0 + undesired_contacts: -0.1 std_root_pos: 0.3 std_root_ori: 0.4 std_body_pos: 0.3 diff --git a/conf/ppo/task/g1_motion_tracking/mujoco.yaml b/conf/ppo/task/g1_motion_tracking/mujoco.yaml index f257cd9e6..e53576b3a 100644 --- a/conf/ppo/task/g1_motion_tracking/mujoco.yaml +++ b/conf/ppo/task/g1_motion_tracking/mujoco.yaml @@ -24,6 +24,7 @@ reward: motion_joint_vel: 0.0 action_rate_l2: -0.1 joint_limit: -10.0 + undesired_contacts: -0.1 std_root_pos: 0.3 std_root_ori: 0.4 std_body_pos: 0.3 diff --git a/src/unilab/assets/motions/g1/run1_subject5_from_csv.npz b/src/unilab/assets/motions/g1/run1_subject5_from_csv.npz new file mode 100644 index 000000000..fa8fa8fa2 Binary files /dev/null and b/src/unilab/assets/motions/g1/run1_subject5_from_csv.npz differ diff --git a/src/unilab/envs/locomotion/g1/base.py b/src/unilab/envs/locomotion/g1/base.py index f72cb419a..3f391c02f 100644 --- a/src/unilab/envs/locomotion/g1/base.py +++ b/src/unilab/envs/locomotion/g1/base.py @@ -13,10 +13,10 @@ @dataclass class NoiseConfig: - level: float = 0.0 - scale_joint_angle: float = 0.02 - scale_joint_vel: float = 0.3 - scale_gyro: float = 0.1 + level: float = 1.0 + scale_joint_angle: float = 0.01 + scale_joint_vel: float = 1.5 + scale_gyro: float = 0.2 scale_gravity: float = 0.05 scale_linvel: float = 0.1 @@ -46,3 +46,14 @@ class G1BaseEnv(LocomotionBaseEnv): _cfg: G1BaseCfg _keyframe_name = "stand" _use_global_dtype = False + + def _obs_noise(self, data: np.ndarray, scale: float) -> np.ndarray: + """Apply per-step uniform observation noise scaled by ``noise_config.level``.""" + noise_cfg = self._cfg.noise_config + if noise_cfg.level > 0.0: + return data + ( + np.random.uniform(-1.0, 1.0, data.shape).astype(data.dtype) + * noise_cfg.level + * scale + ) + return data diff --git a/src/unilab/envs/locomotion/g1/joystick.py b/src/unilab/envs/locomotion/g1/joystick.py index 3aa7d9d01..948358333 100644 --- a/src/unilab/envs/locomotion/g1/joystick.py +++ b/src/unilab/envs/locomotion/g1/joystick.py @@ -217,7 +217,13 @@ def update_state(self, state: NpEnvState) -> NpEnvState: def _compute_obs( self, info: dict, linvel, gyro, gravity, dof_pos, dof_vel ) -> dict[str, np.ndarray]: + noise_cfg = self._cfg.noise_config diff = dof_pos - self.default_angles + gyro = self._obs_noise(gyro, noise_cfg.scale_gyro) + gravity = self._obs_noise(gravity, noise_cfg.scale_gravity) + diff = self._obs_noise(diff, noise_cfg.scale_joint_angle) + dof_vel = self._obs_noise(dof_vel, noise_cfg.scale_joint_vel) + linvel = self._obs_noise(linvel, noise_cfg.scale_linvel) command = info["commands"] last_actions = info.get("current_actions", np.zeros_like(diff)) gait_phase = info.get("gait_phase", np.zeros((self._num_envs, 2), dtype=get_global_dtype())) diff --git a/src/unilab/envs/locomotion/g1/joystick_sac.py b/src/unilab/envs/locomotion/g1/joystick_sac.py index 3699a60a9..3759e77a5 100644 --- a/src/unilab/envs/locomotion/g1/joystick_sac.py +++ b/src/unilab/envs/locomotion/g1/joystick_sac.py @@ -101,7 +101,13 @@ def __init__(self, cfg: G1JoystickSACCfg, num_envs=1, backend_type="mujoco"): def _compute_obs( self, info: dict, linvel, gyro, gravity, dof_pos, dof_vel ) -> dict[str, np.ndarray]: + noise_cfg = self._cfg.noise_config diff = dof_pos - self.default_angles + gyro = self._obs_noise(gyro, noise_cfg.scale_gyro) + gravity = self._obs_noise(gravity, noise_cfg.scale_gravity) + diff = self._obs_noise(diff, noise_cfg.scale_joint_angle) + dof_vel = self._obs_noise(dof_vel, noise_cfg.scale_joint_vel) + linvel = self._obs_noise(linvel, noise_cfg.scale_linvel) command = info["commands"] last_actions = info.get("current_actions", np.zeros_like(diff)) gait_phase = info.get("gait_phase", np.zeros((self._num_envs, 2), dtype=get_global_dtype())) diff --git a/src/unilab/envs/motion_tracking/g1/tracking.py b/src/unilab/envs/motion_tracking/g1/tracking.py index 3225b08bf..90053fbd8 100644 --- a/src/unilab/envs/motion_tracking/g1/tracking.py +++ b/src/unilab/envs/motion_tracking/g1/tracking.py @@ -57,6 +57,7 @@ class RewardConfig: "motion_joint_vel": 0.0, "action_rate_l2": -0.1, "joint_limit": -10.0, + "undesired_contacts": -0.1, } ) # Standard deviations for exponential rewards @@ -119,8 +120,8 @@ class G1MotionTrackingCfg(G1BaseCfg): ASSETS_ROOT_PATH / "motions" / "g1" / "dance1_subject2_part.npz" ) # motion_file: str | list[str] = str(ASSETS_ROOT_PATH / "motions" / "g1" / "gangnam_style.npz") - # motion_file: str | list[str] = str(ASSETS_ROOT_PATH / "motions" / "g1" / "walk1_subject5_from_csv.npz") #LAFAN - # motion_file: str | list[str] = str(ASSETS_ROOT_PATH / "motions" / "g1" / "sprint1_subject4_from_csv.npz") #LAFAN + # motion_file: str | list[str] = str(ASSETS_ROOT_PATH / "motions" / "g1" / "fight1_subject5_from_csv.npz") #LAFAN + # motion_file: str | list[str] = str(ASSETS_ROOT_PATH / "motions" / "g1" / "dance_basic_slide_180_R_loop_001__A322_M.npz") #LAFAN # motion_file: str | list[str] = str(ASSETS_ROOT_PATH / "motions" / "g1" / "playing_violin_R_003__A327_from_csv.npz") #Seed anchor_body_name: str = "torso_link" body_names: tuple[str, ...] = ( @@ -156,6 +157,7 @@ class G1MotionTrackingCfg(G1BaseCfg): "left_wrist_yaw_link", "right_wrist_yaw_link", ) + undesired_contact_z_threshold: float = 0.05 @registry.envcfg("G1MotionTracking") @@ -316,6 +318,13 @@ def __init__(self, cfg: G1MotionTrackingCfg, num_envs=1, backend_type="mujoco"): [cfg.body_names.index(name) for name in cfg.ee_body_names], dtype=np.int32 ) + # Get non-EE body indices for undesired contact penalty + ee_set = set(cfg.ee_body_names) + self.undesired_contact_body_indices = np.array( + [i for i, name in enumerate(cfg.body_names) if name not in ee_set], + dtype=np.int32, + ) + # Load motion data self.motion_loader = MotionLoader(cfg.motion_file, body_indices=motion_body_ids) self.motion_sampler = MotionSampler( @@ -366,6 +375,7 @@ def _init_reward_functions(self): "motion_joint_vel": self._reward_motion_joint_vel, "action_rate_l2": self._reward_action_rate_l2, "joint_limit": self._reward_joint_limit, + "undesired_contacts": self._reward_undesired_contacts, } def update_state(self, state: NpEnvState) -> NpEnvState: @@ -543,6 +553,13 @@ def _compute_obs( joint_pos_rel = dof_pos - self.default_angles last_actions = info.get("current_actions", np.zeros_like(joint_pos_rel)) + # Per-step observation noise on sensor channels + noise_cfg = self._cfg.noise_config + linvel = self._obs_noise(linvel, noise_cfg.scale_linvel) + gyro = self._obs_noise(gyro, noise_cfg.scale_gyro) + joint_pos_rel = self._obs_noise(joint_pos_rel, noise_cfg.scale_joint_angle) + dof_vel = self._obs_noise(dof_vel, noise_cfg.scale_joint_vel) + # Actor observations actor_obs = np.concatenate( [ @@ -703,6 +720,12 @@ def _reward_motion_joint_vel(self, info: dict) -> np.ndarray: np.exp(-error / self._cfg.reward_config.std_joint_vel**2), dtype=get_global_dtype() ) + def _reward_undesired_contacts(self, info: dict) -> np.ndarray: + robot_body_pos_w = info["robot_body_pos_w"] + body_z = robot_body_pos_w[:, self.undesired_contact_body_indices, 2] + is_contact = body_z < self._cfg.undesired_contact_z_threshold + return np.asarray(is_contact.sum(axis=-1), dtype=get_global_dtype()) + def _reward_action_rate_l2(self, info: dict) -> np.ndarray: action_diff = info["current_actions"] - info["last_actions"] return np.asarray(np.sum(np.square(action_diff), axis=1), dtype=get_global_dtype()) diff --git a/tests/envs/test_g1_obs_noise.py b/tests/envs/test_g1_obs_noise.py new file mode 100644 index 000000000..77498eb17 --- /dev/null +++ b/tests/envs/test_g1_obs_noise.py @@ -0,0 +1,78 @@ +"""Tests for G1 per-step observation noise.""" + +from __future__ import annotations + +import numpy as np + +from unilab.envs.locomotion.g1.base import G1BaseCfg, G1BaseEnv, NoiseConfig + + +class _ConcreteG1Env(G1BaseEnv): + """Minimal concrete subclass — only needed to satisfy the ABC.""" + + def update_state(self, state): + raise NotImplementedError + + +def _make_env(level: float) -> G1BaseEnv: + cfg = G1BaseCfg(noise_config=NoiseConfig(level=level)) + env = object.__new__(_ConcreteG1Env) + env._cfg = cfg + return env + + +class TestObsNoise: + def test_noise_applied_when_level_positive(self): + env = _make_env(level=1.0) + data = np.ones((4, 10), dtype=np.float32) + cfg = env._cfg.noise_config + + results = [env._obs_noise(data.copy(), cfg.scale_joint_angle) for _ in range(5)] + # At least one result should differ from the original + assert any(not np.allclose(r, data) for r in results) + + def test_no_noise_when_level_zero(self): + env = _make_env(level=0.0) + data = np.ones((4, 10), dtype=np.float32) + cfg = env._cfg.noise_config + + result = env._obs_noise(data.copy(), cfg.scale_joint_angle) + np.testing.assert_array_equal(result, data) + + def test_noise_bounded_by_level_times_scale(self): + env = _make_env(level=1.0) + data = np.zeros((128, 29), dtype=np.float32) + scale = 0.2 + result = env._obs_noise(data.copy(), scale) + # uniform[-1,1] * 1.0 * 0.2 => bounded by [-0.2, 0.2] + assert np.all(result >= -scale) + assert np.all(result <= scale) + + def test_noise_scales_with_level(self): + env_half = _make_env(level=0.5) + env_full = _make_env(level=1.0) + np.random.seed(42) + data = np.zeros((1024, 10), dtype=np.float32) + scale = 1.0 + + np.random.seed(0) + r_half = env_half._obs_noise(data.copy(), scale) + np.random.seed(0) + r_full = env_full._obs_noise(data.copy(), scale) + + # Same random seed: full-level noise should be exactly 2x half-level noise + np.testing.assert_allclose(r_full, r_half * 2.0) + + def test_noise_preserves_dtype(self): + for dt in [np.float32, np.float64]: + env = _make_env(level=1.0) + data = np.ones((4, 5), dtype=dt) + result = env._obs_noise(data, 0.1) + assert result.dtype == dt + + def test_noise_preserves_shape(self): + env = _make_env(level=1.0) + for shape in [(1, 3), (64, 29), (1024, 10)]: + data = np.zeros(shape, dtype=np.float32) + result = env._obs_noise(data, 0.1) + assert result.shape == shape