diff --git a/docs/sphinx/source/en/4-developer_guide/1-architecture/6-manager_based_api.md b/docs/sphinx/source/en/4-developer_guide/1-architecture/6-manager_based_api.md index 44a29d22e..f30989c77 100644 --- a/docs/sphinx/source/en/4-developer_guide/1-architecture/6-manager_based_api.md +++ b/docs/sphinx/source/en/4-developer_guide/1-architecture/6-manager_based_api.md @@ -24,8 +24,8 @@ historical context. `EntityScene.bind_sensor_data(...)` during construction; their hot path only reads that view and never re-resolves sensor names or XML/model metadata. - `ManagerBasedRlEnv` owns backend materialization exactly once: manager construction - and startup events run first, then `SimBackend.materialize()` completes before any - reset or step can execute. + completes, then `SimBackend.materialize()` and startup events run before any + reset or step. - Explicitly empty configuration may use a Null manager. A requested capability that is unavailable fails at the nearest boundary; it is never skipped, zero-filled, or routed back to a legacy environment. diff --git a/docs/sphinx/source/en/4-developer_guide/2-contracts/4-dr_contract.md b/docs/sphinx/source/en/4-developer_guide/2-contracts/4-dr_contract.md index c31bfa4e3..8b16a5bb7 100644 --- a/docs/sphinx/source/en/4-developer_guide/2-contracts/4-dr_contract.md +++ b/docs/sphinx/source/en/4-developer_guide/2-contracts/4-dr_contract.md @@ -8,11 +8,12 @@ provider protocol and the environment does not carry a DR manager. - **Construction identity:** `env.fixed_model_variants` materializes a final read-only assignment and attaches a UniSim `FixedVariantPlan` to `SceneCfg`. Backends realize it before their first forward and before CUDA graph capture. -- **Startup:** event terms sample one full-width per-environment table before - backend materialization and commit it through the same public - `SimBackend.set_state(..., randomization=...)` boundary as reset DR. The - committed rows become the local reset baseline, so later selected-row resets - preserve unwritten startup values. Startup is opt-in and exchanges continuous +- **Startup:** after manager construction, the backend materializes and event + terms sample one full-width per-environment table. They commit through the + same public `SimBackend.set_state(..., randomization=...)` boundary as reset + DR before read-plan binding, the first reset, or a step. The committed rows + become the local reset baseline, so later selected-row resets preserve + unwritten startup values. Startup is opt-in and exchanges continuous distribution coverage for removing full-model derived-constant refreshes from the steady-state reset path. - **Reset:** event terms write through Entity bindings into diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/6-manager_based_api.md b/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/6-manager_based_api.md index 78caf75ea..a674a5716 100644 --- a/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/6-manager_based_api.md +++ b/docs/sphinx/source/zh_CN/4-developer_guide/1-architecture/6-manager_based_api.md @@ -23,8 +23,8 @@ Torch，并保留现有 `TorchEnvState`、Hydra owner YAML、`SimBackend`、regi `SimBackend` contract，热路径复用缓存 ID/view。 - named-sensor observation term 在构造时通过 `EntityScene.bind_sensor_data(...)` 绑定 backend-owned view；热路径只读该 view，不重复解析 sensor 名称或 XML/model metadata。 -- `ManagerBasedRlEnv` 恰好拥有一次 backend 物化：先完成 manager 构造和 startup event， - 再调用 `SimBackend.materialize()`，任何 reset/step 都不能在物化前执行。 +- `ManagerBasedRlEnv` 恰好拥有一次 backend 物化：先完成 manager 构造，再依次执行 + `SimBackend.materialize()` 和 startup event；任何 reset/step 都不能在物化前执行。 - 用户显式空配置可以使用 Null manager；配置请求但 runtime/backend 不支持的能力必须在 最近边界报错，不能 warning、skip、返回零或回退旧 env。 - 热路径避免明显的重复解析、逐环境 Python 循环、复制和临时分配；进一步优化需要 diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/2-contracts/4-dr_contract.md b/docs/sphinx/source/zh_CN/4-developer_guide/2-contracts/4-dr_contract.md index 839fdf2c5..a761af645 100644 --- a/docs/sphinx/source/zh_CN/4-developer_guide/2-contracts/4-dr_contract.md +++ b/docs/sphinx/source/zh_CN/4-developer_guide/2-contracts/4-dr_contract.md @@ -8,11 +8,12 @@ Manager-Based event term 是 UniLab 唯一的 DR lifecycle。任务 provider 协 - **Construction identity：** `env.fixed_model_variants` 物化最终 read-only assignment，并把 UniSim `FixedVariantPlan` 附到 `SceneCfg`。backend 在 first forward 与 CUDA graph capture 前完成 realization。 -- **Startup：** event term 在 backend materialization 前采样一个 full-width - per-environment 表，并复用公开的 `SimBackend.set_state(..., randomization=...)` - boundary 提交。已提交 rows 成为本地 reset baseline，因此后续 selected-row - reset 会保留未写入的 startup 值。Startup 是显式 opt-in，用连续分布覆盖换取 - steady-state reset 路径中的 full-model derived-constant refresh。 +- **Startup：** Manager 构造完成后，backend 先 materialize，event term 再采样 + 一个 full-width per-environment 表，并在 read-plan 绑定、第一次 reset 或 step + 之前复用公开的 `SimBackend.set_state(..., randomization=...)` boundary 提交。 + 已提交 rows 成为本地 reset baseline，因此后续 selected-row reset 会保留未写入 + 的 startup 值。Startup 是显式 opt-in，用连续分布覆盖换取 steady-state reset + 路径中的 full-model derived-constant refresh。 - **Reset：** event term 通过 Entity binding 写入 `ResetStateTransaction`； transaction 只调用一次 `SimBackend.set_state(..., randomization=...)`。 - **Interval：** event term 通过公开 `SimBackend` contract 使用 backend-owned diff --git a/src/unilab/base/reset_state.py b/src/unilab/base/reset_state.py index 109be9869..7fdc5a3cd 100644 --- a/src/unilab/base/reset_state.py +++ b/src/unilab/base/reset_state.py @@ -257,17 +257,28 @@ def _commit_tensor_event_through_host_bridge( self, host_plan: HostBridgeTransferPlan ) -> dict | None: """Commit owner-staged tensor rows through the packed host bridge.""" - self._last_commit_had_writes = self._tensor_has_writes + host_randomization = self._host_randomization_staged() + self._last_commit_had_writes = self._tensor_has_writes or host_randomization try: - if not self._tensor_has_writes: + if not self._tensor_has_writes and not host_randomization: return None - if self.scene_layout is not None or self._randomization_dirty_masks: + if self.scene_layout is not None: raise NotImplementedError( "packed tensor reset event commit supports scalar qpos/qvel rows only" ) assert self._tensor_rows is not None assert self._tensor_qpos is not None assert self._tensor_qvel is not None + randomization = None + if host_randomization: + capabilities = self._backend.get_tensor_capabilities() + if not capabilities.reset_randomization: + raise NotImplementedError( + "packed tensor reset event commit does not support host-staged reset " + "randomization on backend " + f"'{self._backend.backend_type}'" + ) + randomization = self._tensor_randomization_payload(record=False) rows = self._tensor_rows qpos = self._tensor_qpos.index_select(0, rows).detach() qvel = self._tensor_qvel.index_select(0, rows).detach() @@ -277,7 +288,11 @@ def _commit_tensor_event_through_host_bridge( qpos = qpos.to(device=packed_reset_device, non_blocking=False) qvel = qvel.to(device=packed_reset_device, non_blocking=False) started = time.perf_counter() - result = host_plan.apply_reset(rows, qpos, qvel, randomization=None) + result = host_plan.apply_reset(rows, qpos, qvel, randomization=randomization) + if randomization is not None: + self._record_committed_payload( + self._tensor_rows.detach().cpu().numpy(), randomization + ) self._last_set_state_timing_ms = { "dr_reset_set_state_ms": (time.perf_counter() - started) * 1000.0 } @@ -1910,12 +1925,14 @@ def _commit_tensor_event(self) -> dict | None: finally: self._finish() - def _tensor_randomization_payload(self) -> ResetRandomizationPayload | None: + def _tensor_randomization_payload( + self, *, record: bool = True + ) -> ResetRandomizationPayload | None: """Build the selected-row public DR payload from tensor reset rows.""" assert self._tensor_rows is not None rows_host = self._tensor_rows.detach().cpu().numpy() payload = self._build_randomization_payload(rows_host) - if payload is not None: + if payload is not None and record: self._record_committed_payload(rows_host, payload) return payload diff --git a/src/unilab/envs/manager_based_rl_env.py b/src/unilab/envs/manager_based_rl_env.py index 5ce7eee4a..f1284da49 100644 --- a/src/unilab/envs/manager_based_rl_env.py +++ b/src/unilab/envs/manager_based_rl_env.py @@ -398,16 +398,17 @@ def __init__(self, cfg: ManagerBasedRlEnvCfg, backend: SimBackend, num_envs: int self._configure_action_control() self.set_autoreset(cfg.auto_reset) + self._materialize_backend() if "startup" in self.event_manager.available_modes: # Startup model-field DR uses the same public set_state transaction - # as reset DR. It is intentionally full-width and executes once on - # the cold path, before backend materialization and runtime reads. + # as reset DR. Materialize the backend first: adapters may make the + # selected-row reset boundary and model pools available only there. + # This remains a cold-path commit before read-plan binding/reset. with self._reset_state.scoped(self._all_env_rows): self.event_manager.apply( mode="startup", env_ids=self._all_env_rows if self.event_manager.uses_startup_rows else None, ) - self._materialize_backend() self._compile_tensor_read_plan() self.command_manager.bind_read_phase() self._validate_manager_tensor_runtime() diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index c7ca4807b..36698be8e 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -551,6 +551,110 @@ def count_step(*args, **kwargs): assert "[mujoco tensor motion rollout] OK" in result.stdout +@pytest.mark.parametrize("mode", ["startup", "reset"]) +def test_flashsac_g1_motion_mujoco_tensor_command_model_field_dr_roll_out( + mode: str, +) -> None: + """MuJoCo tensor motion commits startup/reset model-field DR in one boundary.""" + ensure_registries() + _require_mujoco_runtime() + from unilab.base import registry + from unilab.envs import ManagerBasedRlEnv + + _, override = _motion_manager_override("g1_motion_tracking", "mujoco", config_root="flashsac") + override["scene"]["entities"]["robot"]["geom_names"] = [ + f"{side}_foot{index}_collision" for side in ("left", "right") for index in range(1, 8) + ] + events = { + "body_mass": { + "_target_": "unilab.managers.event_manager.EventTermCfg", + "func": "unilab.envs.mdp.randomize_rigid_body_mass", + "mode": mode, + "params": { + "asset_cfg": { + "_target_": "unilab.managers.SceneEntityCfg", + "name": "robot", + "body_names": ["pelvis", "torso_link"], + }, + "mass_distribution_params": [1.0, 1.0], + "operation": "scale", + "recompute_inertia": False, + }, + }, + "body_com": { + "_target_": "unilab.managers.event_manager.EventTermCfg", + "func": "unilab.envs.mdp.randomize_rigid_body_com", + "mode": mode, + "params": { + "asset_cfg": { + "_target_": "unilab.managers.SceneEntityCfg", + "name": "robot", + "body_names": ["pelvis", "torso_link"], + }, + "com_range": {"x": [0.0, 0.0], "y": [0.0, 0.0], "z": [0.0, 0.0]}, + }, + }, + } + if mode == "reset": + events.update( + { + "foot_friction": { + "_target_": "unilab.managers.event_manager.EventTermCfg", + "func": "unilab.envs.mdp.geom_friction", + "mode": "reset", + "params": { + "asset_cfg": { + "_target_": "unilab.managers.SceneEntityCfg", + "name": "robot", + "geom_names": "^(left|right)_foot[1-7]_collision$", + }, + "ranges": [1.0, 1.0], + "operation": "abs", + }, + }, + "pd_gains": { + "_target_": "unilab.managers.event_manager.EventTermCfg", + "func": "unilab.envs.mdp.pd_gains", + "mode": "reset", + "params": { + "asset_cfg": { + "_target_": "unilab.managers.SceneEntityCfg", + "name": "robot", + "actuator_names": ".*", + }, + "kp_range": [1.0, 1.0], + "kd_range": [1.0, 1.0], + "operation": "scale", + }, + }, + } + ) + override["events"] = events + + env = registry.make( + "G1MotionTracking", + num_envs=2, + sim_backend="mujoco", + env_cfg_override=override, + ) + assert isinstance(env, ManagerBasedRlEnv) + try: + command = env.command_manager.get_term("motion") + assert command.tensor_carrier is True + obs, _ = env.reset(seed=1812) + assert all(torch.isfinite(values).all() for values in obs.values()) + for _ in range(3): + state = env.step(torch.zeros((2, 29), dtype=torch.float32)) + assert torch.isfinite(state.reward).all() + assert all(torch.isfinite(values).all() for values in state.obs.values()) + + obs, _ = env.reset(env_indices=torch.tensor([1], dtype=torch.int64)) + assert all(torch.isfinite(values).all() for values in obs.values()) + finally: + env.close() + env._backend.close() + + def test_flashsac_g1_motion_genesis_manager_tensor_command_roll_out() -> None: """The canonical FlashSAC Genesis owner exercises the Manager tensor path.""" ensure_registries() diff --git a/tests/envs/test_manager_based_rl_env.py b/tests/envs/test_manager_based_rl_env.py index 912cd065b..6d4e2e9f9 100644 --- a/tests/envs/test_manager_based_rl_env.py +++ b/tests/envs/test_manager_based_rl_env.py @@ -19,6 +19,7 @@ TensorLifecycleCapabilities, TensorProcessTopology, ) +from unisim.dr.types import DomainRandomizationCapabilities import unilab.envs.manager_based_rl_env as manager_env_module from unilab.assets import ASSETS_ROOT_PATH @@ -53,6 +54,7 @@ RewardTermCfg, TerminationTermCfg, ) +from unilab.managers.scene_entity_config import SceneEntityCfg class _FakeBackend: @@ -87,6 +89,53 @@ def materialize(self) -> None: def get_actuator_names(self) -> tuple[str, ...]: return ("motor",) + def get_geom_names(self) -> tuple[str, ...]: + return ("base_geom",) + + def get_actuator_gains(self) -> tuple[np.ndarray, np.ndarray]: + return np.asarray([20.0]), np.asarray([2.0]) + + def get_body_ids(self, names) -> np.ndarray: + available = ("base",) + return np.asarray([available.index(name) for name in names], dtype=np.int32) + + def get_dr_capabilities(self): + return DomainRandomizationCapabilities( + supported_reset_terms=frozenset(("body_mass", "body_ipos", "geom_friction", "kp", "kd")) + ) + + def get_reset_term_default(self, term: str) -> np.ndarray: + defaults = { + "body_mass": np.asarray([10.0]), + "body_ipos": np.asarray([[0.0, 0.0, 0.0]]), + "geom_friction": np.asarray([[0.5, 0.02, 0.002]]), + "kp": np.asarray([20.0]), + "kd": np.asarray([2.0]), + } + if term not in defaults: + raise NotImplementedError(term) + return np.array(defaults[term], copy=True) + + def get_body_pos_w(self, body_ids) -> np.ndarray: + return np.zeros((self.num_envs, len(body_ids), 3), dtype=np.float32) + + def get_body_quat_w(self, body_ids) -> np.ndarray: + values = np.zeros((self.num_envs, len(body_ids), 4), dtype=np.float32) + values[..., 0] = 1.0 + return values + + def get_body_lin_vel_w(self, body_ids) -> np.ndarray: + return np.zeros((self.num_envs, len(body_ids), 3), dtype=np.float32) + + def get_body_ang_vel_w(self, body_ids) -> np.ndarray: + return np.zeros((self.num_envs, len(body_ids), 3), dtype=np.float32) + + def get_body_lin_vel_b(self, body_ids) -> np.ndarray: + return self.get_body_lin_vel_w(body_ids) + + def get_body_ang_vel_b(self, body_ids) -> np.ndarray: + return self.get_body_ang_vel_w(body_ids) + def get_actuator_ctrl_range(self) -> np.ndarray: return np.array([[-2.0, 2.0]], dtype=np.float32) @@ -114,6 +163,7 @@ def get_tensor_capabilities(self) -> TensorLifecycleCapabilities: data_plane=TensorDataPlane.HOST_BRIDGE, stream_event_ownership="test", torch_devices=("cpu",), + reset_randomization=True, ) def get_public_state_widths(self): @@ -154,6 +204,7 @@ def __init__(self, num_envs: int) -> None: self.init_qvel_calls = 0 self.joint_layout_calls = 0 self.set_state_calls: list[tuple[np.ndarray, np.ndarray, np.ndarray]] = [] + self.randomization_calls: list[Any] = [] def get_default_qpos(self) -> np.ndarray: self.default_qpos_calls += 1 @@ -203,9 +254,9 @@ def set_state( qvel: np.ndarray, randomization=None, ) -> None: - assert randomization is None assert self.materialize_calls == 1 self.set_state_calls.append((env_ids.copy(), qpos.copy(), qvel.copy())) + self.randomization_calls.append(randomization) class _StateBackend(_ResetBackend): @@ -322,6 +373,7 @@ def __init__(self, num_envs: int) -> None: self.selected_reads = 0 self.plan_closes = 0 self.reset_calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = [] + self.reset_randomization_calls: list[Any] = [] self.nq = 4 self.nv = 3 self.sensors = { @@ -331,6 +383,12 @@ def __init__(self, num_envs: int) -> None: ), "track_linvel_w_platform": torch.zeros((num_envs, 3), dtype=torch.float32), "track_angvel_w_platform": torch.zeros((num_envs, 3), dtype=torch.float32), + "track_pos_w_base": torch.zeros((num_envs, 3), dtype=torch.float32), + "track_quat_w_base": torch.tile( + torch.tensor([1.0, 0.0, 0.0, 0.0], dtype=torch.float32), (num_envs, 1) + ), + "track_linvel_w_base": torch.zeros((num_envs, 3), dtype=torch.float32), + "track_angvel_w_base": torch.zeros((num_envs, 3), dtype=torch.float32), "track_pos_w_ball": torch.tensor( [[0.1, 0.2, 1.2], [0.3, 0.4, 1.2]], dtype=torch.float32 ), @@ -353,6 +411,7 @@ def get_tensor_capabilities(self) -> TensorLifecycleCapabilities: data_plane=TensorDataPlane.HOST_BRIDGE, stream_event_ownership="test", torch_devices=("cpu",), + reset_randomization=True, ) def get_state_views(self, fields, device=None) -> dict[str, torch.Tensor]: @@ -373,6 +432,8 @@ def get_body_pos_w(self, body_ids) -> np.ndarray: name = ("base", "platform", "ball")[int(body_id)] if name == "platform": positions[:, row] = self.sensors["track_pos_w_platform"] + elif name == "base": + positions[:, row] = 0.1 elif name == "ball": positions[:, row] = self.sensors["track_pos_w_ball"] return positions.numpy() @@ -442,6 +503,7 @@ def apply_reset(self, env_indices, qpos, qvel, randomization=None): qvel.detach().clone(), ) ) + backend.reset_randomization_calls.append(randomization) return None def close(self) -> None: @@ -1461,7 +1523,7 @@ def wrapped(*args, _name=name, _original=original, **kwargs): assert order == list(names) -def test_backend_materializes_once_after_startup_and_before_runtime() -> None: +def test_backend_materializes_once_before_startup_and_runtime() -> None: cfg = _make_cfg() cfg.events = { "startup": EventTermCfg(func=_startup_event, mode="startup"), @@ -1469,10 +1531,64 @@ def test_backend_materializes_once_after_startup_and_before_runtime() -> None: } env, backend = _make_env(cfg) - assert backend.lifecycle == ["startup", "materialize"] + assert backend.lifecycle == ["materialize", "startup"] assert backend.materialize_calls == 1 - env.reset() + +def test_startup_model_field_dr_commits_after_backend_materialization() -> None: + cfg = _make_cfg() + cfg.scene.entities["robot"] = EntityCfg( + root_body_name="base", + body_names=("base",), + joint_names=("joint",), + actuator_names=("motor",), + ) + cfg.observations = { + "actor": ObservationGroupCfg(terms={"policy": ObservationTermCfg(func=_policy_obs)}), + "value": ObservationGroupCfg(terms={"critic": ObservationTermCfg(func=_critic_obs)}), + } + cfg.events = { + "reset_scene": EventTermCfg(func=mdp.reset_scene_to_default, mode="reset"), + "body_mass": EventTermCfg( + func=mdp.randomize_rigid_body_mass, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("robot", body_names=("base",)), + "mass_distribution_params": (1.5, 1.5), + "operation": "scale", + "recompute_inertia": False, + }, + ), + "body_com": EventTermCfg( + func=mdp.randomize_rigid_body_com, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("robot", body_names=("base",)), + "com_range": {"x": (0.1, 0.1), "z": (-0.2, -0.2)}, + }, + ), + } + backend = _ResetBackend(2) + env = _TestEnv(cfg, cast(SimBackend, backend), 2) + + assert backend.lifecycle == ["materialize"] + assert len(backend.set_state_calls) == 1 + np.testing.assert_array_equal(backend.set_state_calls[0][0], np.arange(2)) + payload = backend.randomization_calls[0] + assert payload is not None + assert payload.body_mass is not None + assert payload.body_ipos is not None + np.testing.assert_allclose(payload.body_mass, [[15.0], [15.0]]) + np.testing.assert_allclose(payload.body_ipos, [[[0.1, 0.0, -0.2]]] * 2) + + env.init_state() + assert backend.lifecycle == ["materialize"] + assert len(backend.set_state_calls) == 2 + replay = backend.randomization_calls[-1] + assert replay is None or replay.body_mass is None + assert replay is None or replay.body_ipos is None + + env.reset(env_indices=torch.tensor([0], dtype=torch.int64)) env.step(torch.zeros((2, 1), dtype=torch.float32)) env.reset() @@ -2321,6 +2437,104 @@ def test_scene_read_plan_pairs_reset_with_selected_packed_transfer() -> None: env.close() +def test_host_bridge_tensor_command_reset_carries_model_field_randomization() -> None: + cfg = _make_cfg(include_optional_managers=False) + cfg.observations = { + "actor": ObservationGroupCfg( + terms={"policy": ObservationTermCfg(func=_tensor_runtime_policy_obs)} + ), + "value": ObservationGroupCfg( + terms={"critic": ObservationTermCfg(func=_tensor_runtime_critic_obs)} + ), + } + cfg.scene.entities["robot"] = EntityCfg( + root_body_name="base", + joint_names=("joint",), + body_names=("base", "platform"), + geom_names=("base_geom",), + actuator_names=("motor",), + ) + cfg.actions = { + "tilt": _TensorBodyActionCfg( + entity_name="robot", + top_body_name="platform", + ball_body_name="base", + ) + } + cfg.commands = {"target": _TensorStateWritingCommandCfg(resampling_time_range=(1.0, 1.0))} + cfg.events = { + "body_mass": EventTermCfg( + func=mdp.randomize_rigid_body_mass, + mode="reset", + params={ + "asset_cfg": SceneEntityCfg("robot", body_names=("base",)), + "mass_distribution_params": (1.5, 1.5), + "operation": "scale", + "recompute_inertia": False, + }, + ), + "body_com": EventTermCfg( + func=mdp.randomize_rigid_body_com, + mode="reset", + params={ + "asset_cfg": SceneEntityCfg("robot", body_names=("base",)), + "com_range": {"x": (0.1, 0.1), "z": (-0.2, -0.2)}, + }, + ), + "foot_friction": EventTermCfg( + func=mdp.geom_friction, + mode="reset", + params={ + "asset_cfg": SceneEntityCfg("robot", geom_names=("base_geom",)), + "ranges": (0.75, 0.75), + "operation": "abs", + }, + ), + "pd_gains": EventTermCfg( + func=mdp.pd_gains, + mode="reset", + params={ + "asset_cfg": SceneEntityCfg("robot", actuator_names=("motor",)), + "kp_range": (30.0, 30.0), + "kd_range": (3.0, 3.0), + "operation": "abs", + }, + ), + } + backend = _ScenePlanBackend(2) + env = _TestEnv(cfg, cast(SimBackend, backend), 2) + try: + plan = env.scene._tensor_read_plan + assert plan is not None and plan.host_plan is not None + assert env.command_manager.uses_tensor_reset_rows() + + env.reset() + + assert backend.set_state_calls == [] + assert len(backend.reset_calls) == 1 + rows, qpos, qvel = backend.reset_calls[0] + torch.testing.assert_close(rows, torch.tensor([0, 1], dtype=torch.int64)) + torch.testing.assert_close(qpos[:, 3], torch.full((2,), 0.75)) + torch.testing.assert_close(qvel[:, 2], torch.full((2,), -0.75)) + payload = backend.reset_randomization_calls[0] + assert payload is not None + assert payload.body_mass is not None + assert payload.body_ipos is not None + assert payload.geom_friction is not None + assert payload.kp is not None + assert payload.kd is not None + np.testing.assert_allclose(payload.body_mass, [[15.0], [15.0]]) + np.testing.assert_allclose(payload.body_ipos, [[[0.1, 0.0, -0.2]]] * 2) + np.testing.assert_allclose(payload.geom_friction, [[[0.75, 0.02, 0.002]]] * 2) + np.testing.assert_allclose(payload.kp, [[30.0], [30.0]]) + np.testing.assert_allclose(payload.kd, [[3.0], [3.0]]) + finally: + if env.scene._tensor_read_plan is not None: + env.scene._tensor_read_plan.close() + env.scene._tensor_read_plan = None + env.close() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a local CUDA device") @pytest.mark.parametrize( ("topology", "data_plane"),