diff --git a/src/unilab/base/torch_env.py b/src/unilab/base/torch_env.py index f3c516448..fd2ba6432 100644 --- a/src/unilab/base/torch_env.py +++ b/src/unilab/base/torch_env.py @@ -238,7 +238,9 @@ def _normalize_reset_indices(self, env_indices: torch.Tensor | None) -> torch.Te def step(self, actions: torch.Tensor) -> TorchEnvState: started = time.perf_counter() cpu_started = _cpu_time() + phase = time.perf_counter() self._validate_action(actions) + action_validate_ms = (time.perf_counter() - phase) * 1000.0 self._bind_tensor_runtime() if self._state is None: self.init_state() @@ -305,6 +307,7 @@ def step(self, actions: torch.Tensor) -> TorchEnvState: if not did_reset: self._clear_reset_done_detail_timing(timing) timing["env_step_total_ms"] = (time.perf_counter() - started) * 1000.0 + timing["action_validate_ms"] = action_validate_ms timing["apply_action_ms"] = apply_action_ms timing["apply_action_cpu_ms"] = apply_action_cpu_ms timing["step_core_ms"] = step_core_ms diff --git a/src/unilab/conf/appo/task/g1_walk_flat/base.yaml b/src/unilab/conf/appo/task/g1_walk_flat/base.yaml index c60ab4940..2095192a3 100644 --- a/src/unilab/conf/appo/task/g1_walk_flat/base.yaml +++ b/src/unilab/conf/appo/task/g1_walk_flat/base.yaml @@ -210,7 +210,7 @@ env: params: max_tilt_deg: 25.0 base_height: - func: unilab.envs.mdp.root_height_below_minimum + func: unilab.tasks.locomotion.g1.manager_terms.g1_base_height_below_minimum params: minimum_height: 0.55 policy_observation_group: policy diff --git a/src/unilab/conf/flashsac/task/g1_walk_flat/base.yaml b/src/unilab/conf/flashsac/task/g1_walk_flat/base.yaml index 61e81e005..991ef3555 100644 --- a/src/unilab/conf/flashsac/task/g1_walk_flat/base.yaml +++ b/src/unilab/conf/flashsac/task/g1_walk_flat/base.yaml @@ -206,7 +206,7 @@ env: params: max_tilt_deg: 65.0 base_height: - func: unilab.envs.mdp.root_height_below_minimum + func: unilab.tasks.locomotion.g1.manager_terms.g1_base_height_below_minimum params: minimum_height: 0.3 curriculum: diff --git a/src/unilab/conf/ppo/task/g1_walk_flat/base.yaml b/src/unilab/conf/ppo/task/g1_walk_flat/base.yaml index 7292ff6c1..e61ded6bf 100644 --- a/src/unilab/conf/ppo/task/g1_walk_flat/base.yaml +++ b/src/unilab/conf/ppo/task/g1_walk_flat/base.yaml @@ -210,7 +210,7 @@ env: params: max_tilt_deg: 25.0 base_height: - func: unilab.envs.mdp.root_height_below_minimum + func: unilab.tasks.locomotion.g1.manager_terms.g1_base_height_below_minimum params: minimum_height: 0.55 policy_observation_group: policy diff --git a/src/unilab/conf/sac/task/g1_walk_flat/base.yaml b/src/unilab/conf/sac/task/g1_walk_flat/base.yaml index 61e81e005..991ef3555 100644 --- a/src/unilab/conf/sac/task/g1_walk_flat/base.yaml +++ b/src/unilab/conf/sac/task/g1_walk_flat/base.yaml @@ -206,7 +206,7 @@ env: params: max_tilt_deg: 65.0 base_height: - func: unilab.envs.mdp.root_height_below_minimum + func: unilab.tasks.locomotion.g1.manager_terms.g1_base_height_below_minimum params: minimum_height: 0.3 curriculum: diff --git a/src/unilab/conf/warpsac/task/g1_walk_flat/base.yaml b/src/unilab/conf/warpsac/task/g1_walk_flat/base.yaml index 61e81e005..991ef3555 100644 --- a/src/unilab/conf/warpsac/task/g1_walk_flat/base.yaml +++ b/src/unilab/conf/warpsac/task/g1_walk_flat/base.yaml @@ -206,7 +206,7 @@ env: params: max_tilt_deg: 65.0 base_height: - func: unilab.envs.mdp.root_height_below_minimum + func: unilab.tasks.locomotion.g1.manager_terms.g1_base_height_below_minimum params: minimum_height: 0.3 curriculum: diff --git a/src/unilab/envs/manager_based_rl_env.py b/src/unilab/envs/manager_based_rl_env.py index 7d901a5f7..7f7f11905 100644 --- a/src/unilab/envs/manager_based_rl_env.py +++ b/src/unilab/envs/manager_based_rl_env.py @@ -72,6 +72,19 @@ _DEVICE_RESIDENT_TENSOR_SENSORS: dict[str, frozenset[str]] = { "isaacgym": frozenset({"pelvis_local_linvel", "torso_gyro"}), + "mjwarp": frozenset( + { + "pelvis_local_linvel", + "torso_gyro", + "torso_upvector", + "left_foot_pos", + "right_foot_pos", + "left_foot_quat", + "right_foot_quat", + *(f"left_foot_contact_{index}" for index in range(4)), + *(f"right_foot_contact_{index}" for index in range(4)), + } + ), } @@ -414,6 +427,7 @@ def _compile_tensor_read_plan(self) -> None: specs.extend(self._action_tensor_read_specs()) specs.extend(self._observation_tensor_read_specs()) specs.extend(self._manager_term_tensor_read_specs()) + specs.extend(self._command_tensor_read_specs()) self.scene._tensor_read_plan = ( self.scene.compile_tensor_reads(self.device, specs) if specs else None ) @@ -530,6 +544,67 @@ def _manager_term_tensor_read_specs(self) -> list[SceneTensorReadSpec]: backend_type=backend_type, ) ) + if manager_name == "termination": + term = manager.get_term_cfg(name).func + specs.extend(self._manager_term_body_read_specs(manager_name, name, term)) + return specs + + @staticmethod + def _manager_term_body_read_specs( + manager_name: str, term_name: str, term: Any + ) -> list[SceneTensorReadSpec]: + body_names = getattr(term, "tensor_body_names", None) + if body_names is None: + return [] + if ( + not isinstance(body_names, (tuple, list)) + or any(not isinstance(body_name, str) or not body_name for body_name in body_names) + or len(set(body_names)) != len(body_names) + ): + raise TypeError( + f"ManagerBasedRlEnv tensor body declaration for {manager_name} term " + f"'{term_name}' must be a unique sequence of body names; got {body_names!r}" + ) + params = getattr(getattr(term, "cfg", None), "params", {}) + entity_name = params.get("entity_name") if isinstance(params, dict) else None + if not isinstance(entity_name, str) or not entity_name: + entity_name = "robot" + return [SceneTensorReadSpec(entity=entity_name, body_names=tuple(body_names))] + + def _command_tensor_read_specs(self) -> list[SceneTensorReadSpec]: + """Collect optional command-term reads for the packed phase.""" + specs: list[SceneTensorReadSpec] = [] + narrow = ( + self._backend.get_tensor_capabilities().execution is TensorExecution.DEVICE_RESIDENT + ) + allowed = _DEVICE_RESIDENT_TENSOR_SENSORS.get( + self._backend.backend_type, + {"pelvis_local_linvel", "torso_gyro", "torso_upvector"}, + ) + for name in self.command_manager.active_terms: + sensor_names = getattr(self.command_manager.get_term(name), "tensor_sensor_names", None) + if sensor_names is None: + continue + if not isinstance(sensor_names, (tuple, list)) or any( + not isinstance(sensor_name, str) or not sensor_name for sensor_name in sensor_names + ): + raise TypeError( + "ManagerBasedRlEnv tensor read declaration for command term " + f"'{name}' must be a sequence of sensor names; got {sensor_names!r}" + ) + names = tuple(sensor_names) + if narrow: + names = tuple(sensor_name for sensor_name in names if sensor_name in allowed) + if not names: + continue + command_cfg = self._cfg.commands.get(name) if self._cfg.commands else None + entity_name = getattr(command_cfg, "entity_name", "robot") if command_cfg else "robot" + specs.append( + SceneTensorReadSpec( + entity=entity_name if isinstance(entity_name, str) and entity_name else "robot", + sensor_names=names, + ) + ) return specs @staticmethod diff --git a/src/unilab/envs/mdp/commands/velocity_command.py b/src/unilab/envs/mdp/commands/velocity_command.py index 66e49f701..9e0ff6aeb 100644 --- a/src/unilab/envs/mdp/commands/velocity_command.py +++ b/src/unilab/envs/mdp/commands/velocity_command.py @@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any, cast import numpy as np +import torch from unilab.dtype_config import get_global_dtype from unilab.managers.command_manager import CommandTerm, CommandTermCfg @@ -72,8 +73,8 @@ def __init__(self, cfg: UniformVelocityCommandCfg, env: ManagerBasedRlEnv): self.is_standing_env = np.zeros(self.num_envs, dtype=np.bool_) self.is_world_env = np.zeros(self.num_envs, dtype=np.bool_) self.is_forward_env = np.zeros(self.num_envs, dtype=np.bool_) - self.metrics["error_vel_xy"] = np.zeros(self.num_envs, dtype=dtype) - self.metrics["error_vel_yaw"] = np.zeros(self.num_envs, dtype=dtype) + self.metrics["error_vel_xy"] = torch.zeros(self.num_envs, device=self._device) + self.metrics["error_vel_yaw"] = torch.zeros(self.num_envs, device=self._device) @staticmethod def _validate_cfg(cfg: UniformVelocityCommandCfg) -> None: @@ -117,19 +118,38 @@ def _validate_cfg(cfg: UniformVelocityCommandCfg) -> None: def command(self) -> np.ndarray: return self.vel_command_b + @property + def tensor_sensor_names(self) -> tuple[str, ...]: + """IMU sensors used by per-step command tracking metrics.""" + return ("pelvis_local_linvel", "torso_gyro") + + def _metric_velocities(self) -> tuple[torch.Tensor, torch.Tensor]: + read_plan = getattr(self._env.scene, "_tensor_read_plan", None) + names = self.tensor_sensor_names + if read_plan is not None and set(names).issubset( + read_plan.sensor_names.get(self.cfg.entity_name, ()) + ): + entity = self._env.scene[self.cfg.entity_name] + views = read_plan.sensor_tensor_views(entity, names).values + return views[names[0]], views[names[1]] + # HOST_BRIDGE owners use their public entity facade. This is an + # explicit carrier choice, not a hidden device transfer. + device = getattr(self._env, "device", torch.device("cpu")) + return ( + torch.as_tensor(self.robot.data.root_link_lin_vel_b, device=device), + torch.as_tensor(self.robot.data.root_link_ang_vel_b, device=device), + ) + def _update_metrics(self, env_ids: np.ndarray | None = None) -> None: del env_ids # Metrics accumulate over all rows on every compute. max_command_steps = self.cfg.resampling_time_range[1] / self._env.step_dt + lin_vel, ang_vel = self._metric_velocities() + command = torch.as_tensor(self.vel_command_b, device=lin_vel.device) self.metrics["error_vel_xy"] += ( - np.linalg.norm( - self.vel_command_b[:, :2] - self.robot.data.root_link_lin_vel_b[:, :2], - axis=-1, - ) - / max_command_steps + torch.linalg.vector_norm(command[:, :2] - lin_vel[:, :2], dim=-1) / max_command_steps ) self.metrics["error_vel_yaw"] += ( - np.abs(self.vel_command_b[:, 2] - self.robot.data.root_link_ang_vel_b[:, 2]) - / max_command_steps + torch.abs(command[:, 2] - ang_vel[:, 2]) / max_command_steps ) def _resample_command(self, env_ids: np.ndarray) -> None: diff --git a/src/unilab/tasks/locomotion/g1/manager_terms.py b/src/unilab/tasks/locomotion/g1/manager_terms.py index 45cce03ac..301ec888d 100644 --- a/src/unilab/tasks/locomotion/g1/manager_terms.py +++ b/src/unilab/tasks/locomotion/g1/manager_terms.py @@ -47,6 +47,14 @@ def active_terms(self) -> list[str]: ... def get_term_cfg(self, term_name: str) -> _RewardTermCfgView: ... + class _EntityDataView(Protocol): + @property + def root_link_pos_w(self) -> np.ndarray: ... + + class _EntityView(Protocol): + @property + def data(self) -> _EntityDataView: ... + class _G1Env(ManagerBasedRlEnv, Protocol): @property def common_step_counter(self) -> int: ... @@ -57,6 +65,12 @@ def reset_buf(self) -> np.ndarray: ... @property def reward_manager(self) -> _RewardManagerView: ... + @property + def cfg(self) -> Any: ... + + @property + def scene(self) -> Any: ... + _DEFAULT_ASSET_CFG = SceneEntityCfg("robot") _LEFT_FOOT_CONTACT_SENSORS = tuple(f"left_foot_contact_{index}" for index in range(4)) @@ -674,6 +688,43 @@ def __call__(self, env: ManagerBasedRlEnv, **params: Any) -> np.ndarray | torch. return np.asarray(tilt > self._max_tilt_rad, dtype=np.bool_) +class g1_base_height_below_minimum(ManagerTermBase): + """Terminate when root height drops below the task threshold. + + Declaring the root body lets the packed device-resident read phase serve + this term without the legacy NumPy body-state facade or a full host cache + refresh. + """ + + _allowed_params = frozenset({"minimum_height"}) + + def __init__(self, cfg: ManagerTermBaseCfg, env: _G1Env): + super().__init__(env) + self._minimum_height = _real(self.name, "minimum_height", cfg.params.get("minimum_height")) + asset_cfg = cfg.params.get("asset_cfg", _DEFAULT_ASSET_CFG) + if not isinstance(asset_cfg, SceneEntityCfg): + raise TypeError(f"{self.name} asset_cfg must be a SceneEntityCfg") + self._entity_name = asset_cfg.name + entity_cfg = env.cfg.scene.entities.get(self._entity_name) if env.cfg.scene else None + root = getattr(entity_cfg, "root_body_name", None) if entity_cfg is not None else None + if not isinstance(root, str) or not root: + raise ValueError(f"{self.name} requires an entity root_body_name") + self._root_body_name = root + + @property + def tensor_body_names(self) -> tuple[str, ...]: + return (self._root_body_name,) + + def __call__(self, env: _G1Env, **params: Any) -> np.ndarray | torch.Tensor: + del params + read_plan = getattr(env.scene, "_tensor_read_plan", None) + if read_plan is not None: + body = read_plan.body_tensor_view(self._entity_name, (self._root_body_name,)) + return body.pos_w[:, 0, 2] < self._minimum_height + robot = env.scene[self._entity_name] + return np.asarray(robot.data.root_link_pos_w[:, 2] < self._minimum_height, dtype=np.bool_) + + class penalty_feet_ori(_SensorTerm): """Penalty for non-flat foot orientations (roll/pitch quaternion rows).""" diff --git a/tests/envs/locomotion/g1/test_g1_owner_contract.py b/tests/envs/locomotion/g1/test_g1_owner_contract.py index 3b23f3596..1875d5824 100644 --- a/tests/envs/locomotion/g1/test_g1_owner_contract.py +++ b/tests/envs/locomotion/g1/test_g1_owner_contract.py @@ -20,7 +20,7 @@ from unilab.base import registry 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.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg from unilab.tasks.locomotion.g1 import manager_terms as g1_terms # CPU-bound on the single-core CI runner; kept in the slow lane (make test-slow). @@ -432,7 +432,7 @@ def test_g1_owner_materializes_complete_plain_manager_cfg( assert list(env_cfg.terminations) == ["time_out", "tilt", "base_height"] assert env_cfg.terminations["time_out"].time_out is True assert env_cfg.terminations["tilt"].func is g1_terms.g1_tilt_exceeded - assert env_cfg.terminations["base_height"].func is mdp.root_height_below_minimum + assert env_cfg.terminations["base_height"].func is g1_terms.g1_base_height_below_minimum assert tuple(name for name, term in env_cfg.events.items() if term is not None) == ( expected_events