diff --git a/agent_context/topics/rl-learning/rl-learning.md b/agent_context/topics/rl-learning/rl-learning.md index 622bd5598..842cef483 100644 --- a/agent_context/topics/rl-learning/rl-learning.md +++ b/agent_context/topics/rl-learning/rl-learning.md @@ -168,6 +168,10 @@ consume observations and write action, log-probability, entropy, and value fields needed by their algorithm. Differentiable policies must expose graph-preserving action sampling. +Trainer throughput counts transitions collected during the current `train()` +call. Restored `global_step` remains a cumulative checkpoint counter and is +excluded from that call's SPS numerator. + For actor-critic policies, `policy.obs_groups.actor` and `.critic` select ordered observation groups. The collector and standard buffer preserve separate `critic_obs` when configured; evaluation applies the same selection. The PPO @@ -184,6 +188,11 @@ DexSim's Motion Policy Evaluator, preserving the original task's reset, step, observation and action path. Supplying an Environment means the adapter owns its camera lifecycle; Kit's default flat-ground camera is not applied to it. +Headless evaluation accepts `--renderer no-render`. When a saved training +configuration uses NoRender, Viewer evaluation defaults to Hybrid; an explicit +native renderer overrides that choice. `--viewer --renderer no-render` is +rejected before creating the simulation. + Tasks opt in through `PolicyViewerCameraCfg` and `get_policy_viewer_target_pose()` (world XYZ + XYZW). The six bundled flat velocity tasks own their presets; `policy_evaluation/_viewer_camera.py` owns diff --git a/agent_context/topics/simulation-system/articulation-adapters.md b/agent_context/topics/simulation-system/articulation-adapters.md index 379a572bc..8c9e83fd3 100644 --- a/agent_context/topics/simulation-system/articulation-adapters.md +++ b/agent_context/topics/simulation-system/articulation-adapters.md @@ -43,6 +43,12 @@ full-batch tensors when partial writes must preserve other rows/DOFs. Avoid DexSim's host-materialized selected-DOF path. Pose conversions follow the [public quaternion contract](simulation-system.md#quaternion-and-pose-convention). +`ArticulationData.fetch_state()` reads joint position/velocity and root +pose/velocities together into existing data buffers. The Scene view reuses +DexSim's batch fetch and converts the root-pose layout once. Returned tensors +are borrowed buffers, not snapshots; clone values that must survive later +reads. Fetch after reset or direct writes rather than caching across them. + `Articulation.set_root_velocity()` writes selected world-frame linear and angular velocities as `(N, 6)` rows. The Scene adapter validates the complete input before writing one selected batch; other environment rows are preserved. diff --git a/agent_context/topics/simulation-system/rendering.md b/agent_context/topics/simulation-system/rendering.md index f11de548c..17d3e966d 100644 --- a/agent_context/topics/simulation-system/rendering.md +++ b/agent_context/topics/simulation-system/rendering.md @@ -4,6 +4,15 @@ Read this for physics/render synchronization, native-window/offscreen image processing and readiness reporting. Return to the [simulation overview](simulation-system.md). +## NoRender initialization + +`RenderCfg(renderer="no-render")` selects `Renderer.NORENDER` and requires +headless mode. The manager skips background, light and visual-material setup, +keeps physical ground, and skips Newton render-state publication. DexSim uses +its existing package and device-free NoRender engine. Native camera and window +operations require a native renderer. Checkpoint evaluation with `--viewer` +defaults to Hybrid when the saved training configuration uses NoRender. + ## Rendering does not advance physics `SimulationManager.render_frame()` owns a read-only consumption phase after diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.cfg.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.cfg.rst index 6ca60f663..8ad718a95 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.cfg.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.cfg.rst @@ -172,3 +172,25 @@ to the split modules. physics_cfg_for_backend physics_backend_from_cfg validate_physics_cfg + +Rigid-body property module +-------------------------- + +The rigid-body configuration types are also available from +``embodichain.lab.sim.cfg.rigid``. They describe mass, collision properties, +materials, and backend-specific overrides. ``NewtonCollisionPropertiesCfg`` +accepts an optional ``priority`` for MuJoCo contact-parameter selection; +``None`` preserves the source value. + +.. currentmodule:: embodichain.lab.sim.cfg.rigid + +.. autosummary:: + + MassPropertiesCfg + DefaultRigidBodyPropertiesCfg + CollisionPropertiesCfg + DefaultCollisionPropertiesCfg + NewtonCollisionPropertiesCfg + RigidBodyMaterialCfg + NewtonRigidBodyMaterialCfg + RigidBodyPhysicsCfg diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.objects.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.objects.rst index 79cbe071d..b4ac9fea8 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.objects.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.objects.rst @@ -100,6 +100,12 @@ Rigid Object Group Articulation ------------ +``robot.body_data.fetch_state()`` reads joint position, joint velocity, root +pose and root velocities together through DexSim's batch state interface. +Each call refreshes the values after stepping or state writes. The returned +tensors reuse the data object's buffers; clone tensors when keeping a past +snapshot. Individual data properties remain available for single-field reads. + .. autoclass:: Articulation :members: :inherited-members: diff --git a/docs/source/overview/sim/sensors/contact_sensor.md b/docs/source/overview/sim/sensors/contact_sensor.md index 696acb89c..bfdef8c72 100644 --- a/docs/source/overview/sim/sensors/contact_sensor.md +++ b/docs/source/overview/sim/sensors/contact_sensor.md @@ -142,3 +142,23 @@ may be left over from an earlier update. Newton MuJoCo-Warp exposes contact forces, so both impulse fields are available. Other supported Newton rigid solvers currently expose contact geometry with zero-valued impulse fields. MuJoCo CPU mode does not expose device contact buffers, and DexUni does not currently publish rigid contacts through `ContactQuery`; those modes are therefore unsupported by this sensor. Default Direct GPU reports static counterparts with actor ID `-1` because its raw contact buffer does not expose their object identity. To monitor a dynamic body or articulation link against arbitrary static geometry, select the dynamic/link object and set `filter_need_both_actor=False`. Default CPU and Newton can identify registered static shapes. + + +## CUDA substep sampling + +When a contact history is registered, `BaseEnv` samples the sensor after every +physics substep. On CUDA, the sensor captures dense-row scattering, overflow +accumulation and all history reductions in one graph. The first sample warms +kernels; capture on the following sample does not advance history, and replay +advances it exactly once. Selected-row reset writes into the same live buffers. +Changing the sampling interval, history threshold or counterpart option, or +registering another history rebuilds this graph. + +Contact fetching and actor metadata refresh still run before sampling. DexSim +owns its snapshot and query graphs; the sensor consumes the current query +buffer and its device-resident count. Results are ordered onto the caller's +Torch stream. CPU sampling uses the same reductions without capture. + +The physics loop, simulation clock and rendering callbacks remain outside the +sensor graph. Window, offscreen and browser visualization retain their existing +update order. diff --git a/docs/source/overview/sim/sim_articulation.md b/docs/source/overview/sim/sim_articulation.md index 6588b497b..24acd5cdb 100644 --- a/docs/source/overview/sim/sim_articulation.md +++ b/docs/source/overview/sim/sim_articulation.md @@ -26,6 +26,34 @@ Articulations are configured using the {class}`~cfg.ArticulationCfg` dataclass. At runtime, call `articulation.set_gravity(...)` to change gravity for every environment or for a selected set of environment indices. +### Combined state writes + +Use `Articulation.set_state` (also inherited by `Robot`) to reset selected +articulations in one Scene batch operation: + +```python +robot.set_state( + env_ids=env_ids, + joint_ids=policy_joint_ids, + root_pose=pose, # (N, 7): environment-local xyz + xyzw + qpos=joint_position, # (N, J), clipped to the selected joint limits + target_qpos=joint_position, + clear_dynamics=True, +) +``` + +All supplied joint fields share the selected row and joint order. Optional +`qvel`, `target_qvel`, `qf`, and `root_velocity` fields override reset defaults. +`clear_dynamics=True` clears velocities, forces, external wrenches and solver +history, and holds the final joint positions where no target is supplied. +Newton/MJWarp requires all articulations in each affected solver world when +clearing solver history. Other environments retain their state. + +The Scene batch validates shapes and selections before writes, then propagates +final joint kinematics once. Native execution errors are reported to the caller. +Single-field joint setters use this same writer; target-only writes do not +recompute kinematics. State changes follow the existing render-frame publication lifecycle. + ### Root velocity writes `Articulation.set_root_velocity(velocity, env_ids=None)` writes root-link velocities @@ -142,6 +170,14 @@ Inspect the resolved properties and test the response after changing solvers. ### Joint Position Limits +When binding a finalized Scene, `ArticulationData` initializes position, +velocity and effort limits through one DexSim `ArticulationBatch` property +read. Each environment retains its own values in public DOF order. The +`joint_stiffness`, `joint_damping`, `joint_friction` and `joint_armature` +properties also use this batch interface and return independent snapshots of +the current model. Retained descriptor edits that require a Scene rebuild +become visible after that rebuild. + Use `qpos_limits` to override the limits defined in the asset file. This is the articulation's effective physical limit in simulation, so it is also the range used when `set_qpos(...)` clamps requested joint positions. diff --git a/docs/source/overview/sim/sim_manager/rendering/configuration.md b/docs/source/overview/sim/sim_manager/rendering/configuration.md index 46fccbb53..dd5a7c981 100644 --- a/docs/source/overview/sim/sim_manager/rendering/configuration.md +++ b/docs/source/overview/sim/sim_manager/rendering/configuration.md @@ -9,7 +9,7 @@ and NRD settings used by | Parameter | Type | Default | Description | | :--- | :--- | :--- | :--- | -| `renderer` | `str` | `"auto"` | Renderer backend: `auto`, `hybrid`, `fast-rt`, or `rt`. | +| `renderer` | `str` | `"auto"` | Renderer backend: `auto`, `no-render`, `hybrid`, `fast-rt`, or `rt`. | | `spp` | `int` | `1` | Samples per pixel for ray-traced rendering. Must be at least `1`. | | `denoising` | `DenoisingCfg` | `DenoisingCfg()` | Independent `window`/`offscreen` choices: `off`, `optix`, `dlss`, or `nrd`. | | `tone_mapping_enabled` | `bool` | `False` | Apply modified Reinhard tone mapping to RGB output. | @@ -23,6 +23,16 @@ segmentation masks, normals, and position buffers remain unchanged. ## Renderer selection +For state-based RL, use `RenderCfg(renderer="no-render")` with +`SimulationManagerCfg(headless=True)`. This selects DexSim's `Renderer.NORENDER` +backend. EmbodiChain skips background, light and visual-material creation, +retains physical ground, and skips Newton render-state publication. +Native cameras and windows require `hybrid`, `fast-rt`, or `rt`. +`headless=True` with a native renderer keeps offscreen rendering available. + +Policy evaluation accepts `--renderer no-render`. For visual evaluation of a +NoRender checkpoint, `--viewer` selects Hybrid by default. + With `renderer="auto"`, EmbodiChain selects a backend from the GPU detected at the configured `gpu_id` when the simulation manager is constructed: diff --git a/embodichain/cli/sim.py b/embodichain/cli/sim.py index fc316f6b0..2bc5c78e5 100644 --- a/embodichain/cli/sim.py +++ b/embodichain/cli/sim.py @@ -62,7 +62,7 @@ def add_sim_args_to_parser(parser: argparse.ArgumentParser) -> None: parser.add_argument( "--renderer", type=str, - choices=["auto", "hybrid", "fast-rt", "rt"], + choices=["auto", "no-render", "hybrid", "fast-rt", "rt"], default="auto", help="Renderer backend; omission preserves the launcher/config default.", ) diff --git a/embodichain/lab/gym/envs/base_env.py b/embodichain/lab/gym/envs/base_env.py index 673dabd32..1a7035cf4 100644 --- a/embodichain/lab/gym/envs/base_env.py +++ b/embodichain/lab/gym/envs/base_env.py @@ -727,8 +727,18 @@ def get_obs(self, **kwargs) -> EnvObs: """ with self._profiler.section("proprio"): + # A Python index list makes each proprioception field upload its + # own indices. Reuse the device index until the selection changes. + selection = (tuple(self.active_joint_ids), self.device) + cached = getattr(self, "_proprio_joint_index_cache", None) + if cached is None or cached[0] != selection: + cached = ( + selection, + torch.tensor(selection[0], dtype=torch.long, device=self.device), + ) + self._proprio_joint_index_cache = cached obs = TensorDict( - dict(robot=self.robot.get_proprioception()[:, self.active_joint_ids]), + dict(robot=self.robot.get_proprioception()[:, cached[1]]), batch_size=[self.num_envs], device=self.device, ) @@ -930,7 +940,13 @@ def reset( with self._profiler.section("reset_objects_state"): self.sim.reset_objects_state( - env_ids=reset_ids, excluded_uids=self._detached_uids_for_reset + env_ids=reset_ids, + # Environment sensors are reset below, including sensors + # supplied by custom _setup_sensors implementations. + excluded_uids=[ + *self._detached_uids_for_reset, + *(sensor.uid for sensor in self.sensors.values()), + ], ) for sensor in self.sensors.values(): @@ -940,7 +956,10 @@ def reset( with self._profiler.section("initialize_episode"): self._initialize_episode(reset_ids, **options) self._reset_physical_objective(reset_ids) - self._elapsed_steps[reset_ids] = 0 + elapsed_ids = torch.as_tensor(reset_ids, dtype=torch.long).to( + device=self._elapsed_steps.device, non_blocking=True + ) + self._elapsed_steps.index_fill_(0, elapsed_ids, 0) with self.sim.render_frame(force_visualization=True): with self._profiler.section("get_obs"): diff --git a/embodichain/lab/gym/envs/embodied_env.py b/embodichain/lab/gym/envs/embodied_env.py index 87a281864..8684877ce 100644 --- a/embodichain/lab/gym/envs/embodied_env.py +++ b/embodichain/lab/gym/envs/embodied_env.py @@ -20,6 +20,7 @@ from math import log from functools import wraps from datetime import datetime +import logging import os import threading import copy @@ -902,7 +903,9 @@ def _update_episode_success_status( dtype=torch.bool, device=self.episode_success_status.device, ) - self.episode_success_status[update_mask] = success[update_mask] + self.episode_success_status.copy_( + torch.where(update_mask, success, self.episode_success_status) + ) def _extend_obs(self, obs: EnvObs, **kwargs) -> EnvObs: if self.observation_manager: @@ -990,7 +993,10 @@ def _reset_physical_objective(self, env_ids: Sequence[int] | torch.Tensor) -> No def _initialize_episode( self, env_ids: Sequence[int] | None = None, **kwargs ) -> None: - logger.log_debug(f"Initializing episode for env_ids: {env_ids}", color="blue") + if logger.logger.isEnabledFor(logging.DEBUG): + logger.log_debug( + f"Initializing episode for env_ids: {env_ids}", color="blue" + ) save_data = kwargs.get("save_data", True) # Determine which environments to process @@ -1090,7 +1096,7 @@ def _initialize_episode( _traj_steps = getattr(self, "_traj_steps", None) if _traj_steps is not None: - _traj_steps[env_ids_to_process] = 0 + _traj_steps.index_fill_(0, env_ids_to_process, 0) # Clear episode buffers only after every recorder has consumed them. if self.rollout_buffer is not None and self._rollout_buffer_mode != "rl": @@ -1115,7 +1121,7 @@ def _initialize_episode( self._demo_active_rollout_start_steps[demo_ids] = 0 self._demo_steps[demo_ids] = 0 - self.episode_success_status[env_ids_to_process] = False + self.episode_success_status.index_fill_(0, env_ids_to_process, False) # Stateful managers reset selected rows before reset-mode events run. action_manager = getattr(self, "action_manager", None) diff --git a/embodichain/lab/gym/envs/managers/actions.py b/embodichain/lab/gym/envs/managers/actions.py index 955400a48..26d4a7215 100644 --- a/embodichain/lab/gym/envs/managers/actions.py +++ b/embodichain/lab/gym/envs/managers/actions.py @@ -429,7 +429,25 @@ def apply_actions(self) -> None: ) def reset(self, env_ids: list[int] | torch.Tensor | None = None) -> None: - super().reset(env_ids) + """Clear selected action history and commands without changing bias. + + Args: + env_ids: Rows to reset. None selects all rows. + """ + buffers = ( + self._raw_actions, + self._previous_raw_actions, + self._processed_actions, + ) + if env_ids is None: + for buffer in buffers: + buffer.zero_() + else: + ids = torch.as_tensor(env_ids, dtype=torch.long).to( + device=self.device, non_blocking=True + ) + for buffer in buffers: + buffer.index_fill_(0, ids, 0) class EefPoseAction(ActionTerm): diff --git a/embodichain/lab/gym/envs/managers/event_manager.py b/embodichain/lab/gym/envs/managers/event_manager.py index e1ae2c75c..f35e4f647 100644 --- a/embodichain/lab/gym/envs/managers/event_manager.py +++ b/embodichain/lab/gym/envs/managers/event_manager.py @@ -152,6 +152,7 @@ def __init__(self, cfg: object, env: EmbodiedEnv): # call the base class (this will parse the functors config) super().__init__(cfg, env) + self._all_env_ids = torch.arange(env.num_envs, device=env.device) def __str__(self) -> str: """Returns: A string representation for event manager.""" @@ -310,6 +311,17 @@ def apply( self._call_event_functor( mode, functor_name, functor_cfg, self._env, None ) + elif not functor_cfg.is_global and functor_cfg.interval_step == 1: + # Every row is due, including rows reset on the last step. + # Reuse the known selection instead of synchronizing a + # CUDA nonzero merely to recover all environment indices. + self._call_event_functor( + mode, + functor_name, + functor_cfg, + self._env, + self._all_env_ids, + ) else: valid_env_ids = ( ( diff --git a/embodichain/lab/gym/envs/managers/randomization/physics.py b/embodichain/lab/gym/envs/managers/randomization/physics.py index a3f2af8dc..bbf446cad 100644 --- a/embodichain/lab/gym/envs/managers/randomization/physics.py +++ b/embodichain/lab/gym/envs/managers/randomization/physics.py @@ -16,6 +16,7 @@ from __future__ import annotations +import math from typing import Sequence from embodichain.lab.gym.envs.managers.manager_base import Functor from embodichain.lab.gym.envs.managers.cfg import FunctorCfg @@ -289,6 +290,12 @@ def __init__(self, cfg: FunctorCfg, env: EmbodiedEnv) -> None: if self._asset is None: raise ValueError(f"Articulation '{entity_cfg.uid}' was not found.") self._interval_range_s = tuple(cfg.params["interval_range_s"]) + # A conservative host bound avoids CUDA nonzero on steps where no row + # can be due. Leave one step of slack for float32 interval rounding. + self._minimum_interval_steps = max( + 1, math.floor(float(self._interval_range_s[0]) / env.step_dt) - 1 + ) + self._steps_until_check = self._minimum_interval_steps velocity_range = cfg.params["velocity_range"] self._velocity_lower = torch.tensor( [ @@ -319,6 +326,9 @@ def _resample_interval(self, env_ids: torch.Tensor) -> None: self._steps_remaining[env_ids] = ( torch.ceil(seconds / self._env.step_dt).to(torch.long).clamp_min(1) ) + self._steps_until_check = min( + self._steps_until_check, self._minimum_interval_steps + ) def reset(self, env_ids: Sequence[int] | None = None) -> None: """Resample disturbance timers for selected environments. @@ -331,6 +341,8 @@ def reset(self, env_ids: Sequence[int] | None = None) -> None: if env_ids is None else torch.as_tensor(env_ids, dtype=torch.long, device=self._env.device) ) + if env_ids is None: + self._steps_until_check = self._minimum_interval_steps self._resample_interval(ids) def __call__( @@ -356,14 +368,21 @@ def __call__( if env_ids is None else env_ids.to(device=env.device, dtype=torch.long) ) + if ids.numel() == 0: + return self._steps_remaining[ids] -= 1 - due_ids = ids[self._steps_remaining[ids] <= 0] - if len(due_ids) == 0: + self._steps_until_check -= 1 + if self._steps_until_check > 0: return - random = torch.rand((len(due_ids), 6), device=env.device) - disturbance = self._velocity_lower + random * ( - self._velocity_upper - self._velocity_lower - ) - velocity = self._asset.body_data.root_vel[due_ids] + disturbance - self._asset.set_root_velocity(velocity, env_ids=due_ids) - self._resample_interval(due_ids) + due_ids = ids[self._steps_remaining[ids] <= 0] + if len(due_ids) > 0: + random = torch.rand((len(due_ids), 6), device=env.device) + disturbance = self._velocity_lower + random * ( + self._velocity_upper - self._velocity_lower + ) + velocity = self._asset.body_data.root_vel[due_ids] + disturbance + self._asset.set_root_velocity(velocity, env_ids=due_ids) + self._resample_interval(due_ids) + # All rows bound future calls, including calls advancing a subset. + # Partial resets can only reduce this bound to the minimum interval. + self._steps_until_check = int(self._steps_remaining.min().item()) diff --git a/embodichain/lab/sim/_startup_summary.py b/embodichain/lab/sim/_startup_summary.py index fce8f850a..f6ac636e6 100644 --- a/embodichain/lab/sim/_startup_summary.py +++ b/embodichain/lab/sim/_startup_summary.py @@ -58,6 +58,7 @@ def simulation_rows(sim: SimulationManager) -> list[Row]: cfg = sim.sim_config physics = cfg.physics_cfg render = cfg.render_cfg + native_rendering = render.renderer != "no-render" gpu_name = getattr(sim, "_render_device_name", None) or "name unavailable" gpu = f"{gpu_name} · GPU {cfg.gpu_id}" rows = [ @@ -73,16 +74,24 @@ def simulation_rows(sim: SimulationManager) -> list[Row]: f"{sim.num_envs} · spacing {cfg.arena_space:g} m", ), ("Rendering", "Renderer", _selection(sim._requested_renderer, render.renderer)), - ("Rendering", "Graphics API", "Vulkan"), - ("Rendering", "Render GPU", gpu), + ("Rendering", "Graphics API", "Vulkan" if native_rendering else "OFF"), + ("Rendering", "Render GPU", gpu if native_rendering else "OFF"), ("Rendering", "Native window", "OPEN" if sim.is_window_opened else "CLOSED"), ( "Rendering", "Browser viewer", cfg.visualization.backend if cfg.visualization.backend != "none" else "OFF", ), - ("Rendering", "Viewer resolution", f"{cfg.width} × {cfg.height}"), - ("Rendering", "Tone mapping", _switch(render.tone_mapping_enabled)), + ( + "Rendering", + "Viewer resolution", + f"{cfg.width} × {cfg.height}" if native_rendering else "OFF", + ), + ( + "Rendering", + "Tone mapping", + _switch(native_rendering and render.tone_mapping_enabled), + ), ( "Physics", "Backend", @@ -113,7 +122,7 @@ def simulation_rows(sim: SimulationManager) -> list[Row]: ("Physics", "Gradients", _switch(physics.requires_grad)), ] ) - if cfg.startup_summary == "full": + if cfg.startup_summary == "full" and native_rendering: rows.extend( [ ( @@ -130,6 +139,7 @@ def simulation_rows(sim: SimulationManager) -> list[Row]: ), ] ) + if cfg.startup_summary == "full": if isinstance(physics, DefaultPhysicsCfg): rows.extend( [ @@ -172,11 +182,7 @@ def simulation_rows(sim: SimulationManager) -> list[Row]: ( "Physics detail", "Broad phase", - str( - collision.broad_phase - or physics.broad_phase - or "backend default" - ), + str(collision.broad_phase or "backend default"), ), ( "Physics detail", diff --git a/embodichain/lab/sim/cfg/__init__.py b/embodichain/lab/sim/cfg/__init__.py index 0752b5251..c4ddd4fef 100644 --- a/embodichain/lab/sim/cfg/__init__.py +++ b/embodichain/lab/sim/cfg/__init__.py @@ -79,7 +79,7 @@ from .viewer import MarkerCfg, WindowCameraPoseCfg, WindowRecordCfg # The renderer selection code intentionally mutates this package-level value. -DEFAULT_RENDERER: Literal["auto", "hybrid", "fast-rt", "rt"] = "auto" +DEFAULT_RENDERER: Literal["auto", "no-render", "hybrid", "fast-rt", "rt"] = "auto" # Robot imports are kept last because SolverCfg discovery imports simulation # modules that themselves rely on the public facade above. diff --git a/embodichain/lab/sim/cfg/rendering.py b/embodichain/lab/sim/cfg/rendering.py index 0a9421a66..4985d31b7 100644 --- a/embodichain/lab/sim/cfg/rendering.py +++ b/embodichain/lab/sim/cfg/rendering.py @@ -389,10 +389,12 @@ def to_dexsim_cfg(self, window_width: int, window_height: int) -> dexsim.DLSSCon @configclass class RenderCfg: - renderer: Literal["auto", "hybrid", "fast-rt", "rt"] = "auto" - """Renderer backend to use for the simulation. Options are 'auto', 'hybrid', 'fast-rt', and 'rt'. + renderer: Literal["auto", "no-render", "hybrid", "fast-rt", "rt"] = "auto" + """Renderer backend to use for the simulation. Note: + - 'no-render' selects DexSim's NoRender backend and requires headless mode. + Native camera sensors require 'hybrid', 'fast-rt', or 'rt'. - 'auto' selects a default renderer based on the detected GPU: RTX-series cards use 'hybrid', while datacenter cards (A100/A800, H100/H800/H200/H20) use 'fast-rt'. If no CUDA device is available or the GPU is unknown, it falls back to 'hybrid'. @@ -431,7 +433,9 @@ def __post_init__(self) -> None: def to_dexsim_flags(self) -> Renderer: """Convert the renderer name to DexSim's renderer enum.""" - if self.renderer == "hybrid": + if self.renderer == "no-render": + return Renderer.NORENDER + elif self.renderer == "hybrid": return Renderer.HYBRID elif self.renderer == "fast-rt": return Renderer.FASTRT @@ -447,7 +451,7 @@ def to_dexsim_flags(self) -> Renderer: return Renderer.HYBRID else: logger.log_error( - f"Invalid renderer type '{self.renderer}' specified. Must be one of 'auto', 'hybrid', 'fast-rt', or 'rt'." + f"Invalid renderer type '{self.renderer}' specified. Must be one of 'auto', 'no-render', 'hybrid', 'fast-rt', or 'rt'." ) def apply_to_dexsim_config(self, world_config: dexsim.WorldConfig) -> None: @@ -457,6 +461,8 @@ def apply_to_dexsim_config(self, world_config: dexsim.WorldConfig) -> None: world_config: DexSim world configuration to update in place. """ world_config.renderer = self.to_dexsim_flags() + if self.renderer == "no-render": + return window_mode, offscreen_mode = self.denoising.to_dexsim_modes() set_rt_render_modes = getattr(world_config, "set_rt_render_modes", None) if callable(set_rt_render_modes): diff --git a/embodichain/lab/sim/cfg/rigid.py b/embodichain/lab/sim/cfg/rigid.py index 2f18f6f2f..2d32f5b4a 100644 --- a/embodichain/lab/sim/cfg/rigid.py +++ b/embodichain/lab/sim/cfg/rigid.py @@ -28,6 +28,17 @@ from embodichain.utils import configclass from embodichain.utils.math import quat_wxyz_to_xyzw, quat_xyzw_to_wxyz +__all__ = [ + "MassPropertiesCfg", + "DefaultRigidBodyPropertiesCfg", + "CollisionPropertiesCfg", + "DefaultCollisionPropertiesCfg", + "NewtonCollisionPropertiesCfg", + "RigidBodyMaterialCfg", + "NewtonRigidBodyMaterialCfg", + "RigidBodyPhysicsCfg", +] + @configclass class MassPropertiesCfg: @@ -223,6 +234,9 @@ class NewtonCollisionPropertiesCfg(CollisionPropertiesCfg): `_. """ + priority: int | None = None + """MuJoCo contact-parameter priority; None preserves the source value.""" + condim: int | None = None """MuJoCo contact dimension: 1, 3, 4, or 6. diff --git a/embodichain/lab/sim/objects/articulation.py b/embodichain/lab/sim/objects/articulation.py index 8b55e1ae8..da7fbe67f 100644 --- a/embodichain/lab/sim/objects/articulation.py +++ b/embodichain/lab/sim/objects/articulation.py @@ -32,6 +32,7 @@ ArticulationFlag, DriveType, ) + from dexsim.engine import MaterialInst from embodichain.lab.sim import VisualMaterialInst, VisualMaterial, ReuseSegmentState @@ -284,21 +285,26 @@ def __init__( self._qf = torch.zeros( (self.num_instances, self.dof), dtype=torch.float32, device=self.device ) - self._qpos_limits = torch.as_tensor( - np.array([entity.get_joint_position_limits() for entity in self.entities]), + self._qpos_limits = torch.empty( + (self.num_instances, self.dof, 2), dtype=torch.float32, device=self.device, ) - self._qvel_limits = torch.as_tensor( - np.array([entity.get_joint_velocity_limit() for entity in self.entities]), + self._qvel_limits = torch.empty( + (self.num_instances, self.dof), dtype=torch.float32, device=self.device, ) - self._qf_limits = torch.as_tensor( - np.array([entity.get_joint_effort_limit() for entity in self.entities]), + self._qf_limits = torch.empty( + (self.num_instances, self.dof), dtype=torch.float32, device=self.device, ) + self.articulation_view.fetch_joint_properties( + position_limits=self._qpos_limits, + velocity_limit=self._qvel_limits, + effort_limit=self._qf_limits, + ) @property def is_newton_backend(self) -> bool: @@ -308,6 +314,34 @@ def is_newton_backend(self) -> bool: def is_ready(self) -> bool: return self.articulation_view.is_ready + def fetch_state(self) -> dict[str, torch.Tensor]: + """Read current joint and root state together into reusable buffers. + + Returns: + A mapping with ``qpos``, ``qvel``, ``root_pose``, ``root_lin_vel``, + and ``root_ang_vel``. Poses use ``xyz + xyzw`` and root velocities + are in the world frame. Tensor storage belongs to this data object + and is overwritten by subsequent reads of the same fields. Clone + a returned tensor when retaining a historical snapshot. + + This call always fetches live data; it does not reuse a previous step's + values. Scene views delegate to DexSim's batched state reader. + """ + self.articulation_view.fetch_state( + self._qpos, + self._qvel, + self._root_pose, + self._root_lin_vel, + self._root_ang_vel, + ) + return { + "qpos": self._qpos, + "qvel": self._qvel, + "root_pose": self._root_pose, + "root_lin_vel": self._root_lin_vel, + "root_ang_vel": self._root_ang_vel, + } + @property def root_pose(self) -> torch.Tensor: """Get the root pose of the articulation. @@ -571,6 +605,15 @@ def capture_default_physical_properties( self._default_inertia = inertia.to(self.device, dtype=torch.float32).clone() self._default_com_pose = com_pose.to(self.device, dtype=torch.float32).clone() + def _read_joint_property(self, field_name: str) -> torch.Tensor: + values = torch.empty( + (self.num_instances, self.dof), + dtype=torch.float32, + device=self.device, + ) + self.articulation_view.fetch_joint_properties(**{field_name: values}) + return values + @property def joint_stiffness(self) -> torch.Tensor: """Get the joint stiffness of the articulation. @@ -578,13 +621,7 @@ def joint_stiffness(self) -> torch.Tensor: Returns: torch.Tensor: The joint stiffness of the articulation with shape (N, dof). """ - return torch.as_tensor( - np.array( - [self._entity_drive_properties(entity)[0] for entity in self.entities] - ), - dtype=torch.float32, - device=self.device, - ) + return self._read_joint_property("stiffness") @property def joint_damping(self) -> torch.Tensor: @@ -593,13 +630,7 @@ def joint_damping(self) -> torch.Tensor: Returns: torch.Tensor: The joint damping of the articulation with shape (N, dof). """ - return torch.as_tensor( - np.array( - [self._entity_drive_properties(entity)[1] for entity in self.entities] - ), - dtype=torch.float32, - device=self.device, - ) + return self._read_joint_property("damping") @property def joint_friction(self) -> torch.Tensor: @@ -608,13 +639,7 @@ def joint_friction(self) -> torch.Tensor: Returns: torch.Tensor: The joint friction of the articulation with shape (N, dof). """ - return torch.as_tensor( - np.array( - [self._entity_drive_properties(entity)[4] for entity in self.entities] - ), - dtype=torch.float32, - device=self.device, - ) + return self._read_joint_property("friction") @property def joint_armature(self) -> torch.Tensor: @@ -623,13 +648,7 @@ def joint_armature(self) -> torch.Tensor: Returns: torch.Tensor: The joint armature of the articulation with shape (N, dof). """ - return torch.as_tensor( - np.array( - [self._entity_drive_properties(entity)[5] for entity in self.entities] - ), - dtype=torch.float32, - device=self.device, - ) + return self._read_joint_property("armature") @property def qpos_limits(self) -> torch.Tensor: @@ -1796,6 +1815,96 @@ def set_qf_limits( f"set_joint_effort_limit failed for envs {failed_envs} and joint_ids {joint_ids_np.tolist()}." ) + def set_state( + self, + *, + env_ids: Sequence[int] | torch.Tensor | None = None, + joint_ids: Sequence[int] | torch.Tensor | None = None, + root_pose: torch.Tensor | None = None, + qpos: torch.Tensor | None = None, + target_qpos: torch.Tensor | None = None, + qvel: torch.Tensor | None = None, + target_qvel: torch.Tensor | None = None, + qf: torch.Tensor | None = None, + root_velocity: torch.Tensor | None = None, + clear_dynamics: bool = False, + ) -> None: + """Write selected physical state in one backend operation. + + Args: + env_ids: Environment rows; None selects all environments. + joint_ids: Joint columns for all supplied joint tensors. + root_pose: Environment-local xyz + xyzw, shape (N, 7). + qpos: Current joint positions, shape (N, J), clipped to limits. + target_qpos: Position targets, shape (N, J), clipped to limits. + qvel: Current joint velocities, shape (N, J). + target_qvel: Velocity targets, shape (N, J). + qf: Generalized joint efforts, shape (N, J). + root_velocity: Root world linear + angular velocity, shape (N, 6). + clear_dynamics: Clear velocities, forces and solver history for + selected environments and hold their final joint positions. + Explicit fields override these defaults. + + Unspecified positions and unselected environments retain their values. + Scene backends validate all field shapes before changing physics state. + Joint/root positions are propagated after the final fields are written. + + Raises: + ValueError: A supplied field has an incompatible shape. + RuntimeError: The backend cannot write the selected state. + """ + local_env_ids = self._resolve_env_ids(env_ids) + local_joint_ids = self._resolve_joint_ids(joint_ids) + fields: dict[str, torch.Tensor] = {} + limits = None + for name, value in ( + ("qpos", qpos), + ("target_qpos", target_qpos), + ("qvel", qvel), + ("target_qvel", target_qvel), + ("qf", qf), + ): + if value is None: + continue + value = torch.as_tensor(value, dtype=torch.float32, device=self.device) + if value.ndim == 1: + value = value.unsqueeze(0) + expected = (len(local_env_ids), len(local_joint_ids)) + if tuple(value.shape) != expected: + raise ValueError( + f"Expected {name} shape {expected}, got {tuple(value.shape)}." + ) + if name in ("qpos", "target_qpos"): + if limits is None: + limits = self.body_data.qpos_limits[local_env_ids][ + :, local_joint_ids + ] + value = value.clamp(limits[..., 0], limits[..., 1]) + fields[name] = value + for name, value, width in ( + ("root_pose", root_pose, 7), + ("root_velocity", root_velocity, 6), + ): + if value is not None: + value = torch.as_tensor(value, device=self.device, dtype=torch.float32) + expected = (len(local_env_ids), width) + if tuple(value.shape) != expected: + raise ValueError( + f"Expected {name} shape {expected}, got {tuple(value.shape)}." + ) + fields[name] = value + self._data.articulation_view.apply_state( + None if env_ids is None else local_env_ids, + joint_ids, + clear_dynamics=clear_dynamics, + **fields, + ) + for name, velocity in (("target_qpos", False), ("target_qvel", True)): + if name in fields: + self._stabilize_newton_mimic_target_write( + fields[name], local_env_ids, local_joint_ids, velocity=velocity + ) + def set_qpos( self, qpos: torch.Tensor, @@ -1814,46 +1923,11 @@ def set_qpos( Raises: ValueError: If the length of `env_ids` does not match the length of `qpos`. """ - # TODO: Refactor this part to use a more generic and extensible approach, - # such as a class decorator that can automatically convert ndarray to torch.Tensor - # and handle dimension padding for specified member functions. - # This will make the codebase cleaner and reduce repetitive type checks/conversions. - # (e.g., support specifying which methods should be decorated for auto-conversion.) - if not isinstance(qpos, torch.Tensor): - qpos = torch.as_tensor(qpos, dtype=torch.float32, device=self.device) - else: - qpos = qpos.to(device=self.device, dtype=torch.float32) - - local_joint_ids = self._resolve_joint_ids(joint_ids) - local_env_ids = self._resolve_env_ids(env_ids) - - # Make sure qpos is 2D tensor - if qpos.dim() == 1: - qpos = qpos.unsqueeze(0) - - if len(local_env_ids) != len(qpos): - logger.log_error( - f"Length of env_ids {len(local_env_ids)} does not match qpos length {len(qpos)}. " - f"env_ids: {local_env_ids}, qpos.shape: {qpos.shape}" - ) - - selected_limits = self.body_data.qpos_limits[local_env_ids][ - :, local_joint_ids, : - ] - qpos = qpos.clamp(selected_limits[..., 0], selected_limits[..., 1]) - self._data.articulation_view.apply_qpos( - qpos, - None if env_ids is None else local_env_ids, - None if joint_ids is None else local_joint_ids, - target=target, + self.set_state( + env_ids=env_ids, + joint_ids=joint_ids, + **{"target_qpos" if target else "qpos": qpos}, ) - if target: - self._stabilize_newton_mimic_target_write( - qpos, - local_env_ids, - local_joint_ids, - velocity=False, - ) def get_qvel(self, target: bool = False) -> torch.Tensor: """Get the current velocities (qvel) or target velocities (target_qvel) of the articulation. @@ -1903,36 +1977,11 @@ def set_qvel( Raises: ValueError: If the length of `env_ids` does not match the length of `qvel`. """ - local_env_ids = self._resolve_env_ids(env_ids) - - if not isinstance(qvel, torch.Tensor): - qvel = torch.as_tensor(qvel, dtype=torch.float32, device=self.device) - else: - qvel = qvel.to(device=self.device, dtype=torch.float32) - - if qvel.dim() == 1: - qvel = qvel.unsqueeze(0) - - if len(local_env_ids) != len(qvel): - logger.log_error( - f"Length of env_ids {len(local_env_ids)} does not match qvel length {len(qvel)}." - ) - - local_joint_ids = self._resolve_joint_ids(joint_ids) - - self._data.articulation_view.apply_qvel( - qvel, - None if env_ids is None else local_env_ids, - None if joint_ids is None else local_joint_ids, - target=target, + self.set_state( + env_ids=env_ids, + joint_ids=joint_ids, + **{"target_qvel" if target else "qvel": qvel}, ) - if target: - self._stabilize_newton_mimic_target_write( - qvel, - local_env_ids, - local_joint_ids, - velocity=True, - ) def set_qf( self, @@ -1947,28 +1996,7 @@ def set_qf( joint_ids (Sequence[int] | None, optional): Joint indices to apply the efforts. If None, applies to all joints. env_ids (Sequence[int] | None, optional): Environment indices. Defaults to all indices. """ - local_env_ids = self._resolve_env_ids(env_ids) - - if not isinstance(qf, torch.Tensor): - qf = torch.as_tensor(qf, dtype=torch.float32, device=self.device) - else: - qf = qf.to(device=self.device, dtype=torch.float32) - - if qf.dim() == 1: - qf = qf.unsqueeze(0) - - if len(local_env_ids) != len(qf): - logger.log_error( - f"Length of env_ids {len(local_env_ids)} does not match qf length {len(qf)}." - ) - - local_joint_ids = self._resolve_joint_ids(joint_ids) - - self._data.articulation_view.apply_qf( - qf, - None if env_ids is None else local_env_ids, - None if joint_ids is None else local_joint_ids, - ) + self.set_state(env_ids=env_ids, joint_ids=joint_ids, qf=qf) def get_qf(self) -> torch.Tensor: """Get the current generalized efforts (qf) of the articulation. diff --git a/embodichain/lab/sim/objects/backends/base.py b/embodichain/lab/sim/objects/backends/base.py index b76e60c16..d7fc258de 100644 --- a/embodichain/lab/sim/objects/backends/base.py +++ b/embodichain/lab/sim/objects/backends/base.py @@ -310,6 +310,60 @@ def fetch_link_physical_properties( """ raise NotImplementedError("This view does not expose batch link properties.") + def fetch_joint_properties( + self, + *, + position_limits: torch.Tensor | None = None, + velocity_limit: torch.Tensor | None = None, + effort_limit: torch.Tensor | None = None, + stiffness: torch.Tensor | None = None, + damping: torch.Tensor | None = None, + friction: torch.Tensor | None = None, + armature: torch.Tensor | None = None, + ) -> None: + """Read current joint parameters in public DOF order. + + Args: + position_limits: Lower/upper output, shape ``(N, dof, 2)``. + velocity_limit: Velocity-limit output, shape ``(N, dof)``. + effort_limit: Effort-limit output, shape ``(N, dof)``. + stiffness: Drive-stiffness output, shape ``(N, dof)``. + damping: Drive-damping output, shape ``(N, dof)``. + friction: Joint-friction output, shape ``(N, dof)``. + armature: Joint-armature output, shape ``(N, dof)``. + + Outputs are caller-owned. Omitted fields are not read; at least one + output is required. Values follow the backend's current model, with + retained descriptor changes becoming visible after Scene rebuild. + + Raises: + NotImplementedError: If this view has no batch property reader. + """ + raise NotImplementedError("This view does not expose batch joint properties.") + + def fetch_state( + self, + qpos: torch.Tensor, + qvel: torch.Tensor, + root_pose: torch.Tensor, + root_lin_vel: torch.Tensor, + root_ang_vel: torch.Tensor, + ) -> None: + """Read joint and root state into caller-owned buffers. + + Args: + qpos: Joint positions, shape ``(N, dof)``. + qvel: Joint velocities, shape ``(N, dof)``. + root_pose: Root poses, shape ``(N, 7)``, in ``xyz + xyzw`` order. + root_lin_vel: Root world-frame linear velocities, shape ``(N, 3)``. + root_ang_vel: Root world-frame angular velocities, shape ``(N, 3)``. + """ + self.fetch_qpos(qpos) + self.fetch_qvel(qvel) + self.fetch_root_pose(root_pose) + self.fetch_root_linear_velocity(root_lin_vel) + self.fetch_root_angular_velocity(root_ang_vel) + @abstractmethod def fetch_root_pose(self, data: torch.Tensor) -> torch.Tensor: """Fetch root poses into ``data`` and return a view/result tensor.""" @@ -389,6 +443,38 @@ def apply_root_velocity( """ ... + @abstractmethod + def apply_state( + self, + env_ids: Sequence[int] | torch.Tensor | None = None, + joint_ids: Sequence[int] | torch.Tensor | None = None, + *, + root_pose: torch.Tensor | None = None, + qpos: torch.Tensor | None = None, + target_qpos: torch.Tensor | None = None, + qvel: torch.Tensor | None = None, + target_qvel: torch.Tensor | None = None, + qf: torch.Tensor | None = None, + root_velocity: torch.Tensor | None = None, + clear_dynamics: bool = False, + ) -> None: + """Submit selected final physical state without intermediate propagation. + + Args: + env_ids: Environment rows; None selects all rows. + joint_ids: Joint columns; None selects all joints. + root_pose: Environment-local xyz + xyzw poses, shape ``(N, 7)``. + qpos: Joint positions, shape ``(N, J)``. + target_qpos: Joint position targets, shape ``(N, J)``. + qvel: Joint velocities, shape ``(N, J)``. + target_qvel: Joint velocity targets, shape ``(N, J)``. + qf: Joint efforts, shape ``(N, J)``. + root_velocity: World linear + angular velocities, shape ``(N, 6)``. + clear_dynamics: Clear selected rows' dynamic state and hold their + final joint positions. Explicit fields override these defaults. + """ + ... + @abstractmethod def apply_qpos( self, diff --git a/embodichain/lab/sim/objects/backends/scene.py b/embodichain/lab/sim/objects/backends/scene.py index a27cafdff..20c0d65da 100644 --- a/embodichain/lab/sim/objects/backends/scene.py +++ b/embodichain/lab/sim/objects/backends/scene.py @@ -426,21 +426,22 @@ def __init__( self.scene = scene self._validate_homogeneous_layout() layouts = tuple(self.batch.joint_layouts_per_articulation) + self._joint_dof_columns_host = tuple( + [layout.dof_start for layout in layouts[0]] + if layouts + else list(range(self.batch.dof_width)) + ) self._joint_dof_columns = torch.as_tensor( - ( - [layout.dof_start for layout in layouts[0]] - if layouts - else list(range(self.batch.dof_width)) - ), + self._joint_dof_columns_host, dtype=torch.long, device=device, ) self._articulation_ids = torch.arange( len(batch), dtype=torch.int32, device=device ) - self._joint_apply_scratch: dict[str, torch.Tensor] = {} self._link_com_scratch: torch.Tensor | None = None self._root_velocity_snapshot: torch.Tensor | None = None + self._root_pose_scratch: torch.Tensor | None = None def _validate_homogeneous_layout(self) -> None: """Require the uniform topology promised by one EC Articulation.""" @@ -541,6 +542,63 @@ def fetch_link_physical_properties( com_pose[..., :3].copy_(scratch[..., 4:]) com_pose[..., 3:].copy_(scratch[..., :4]) + def fetch_joint_properties( + self, + *, + position_limits: torch.Tensor | None = None, + velocity_limit: torch.Tensor | None = None, + effort_limit: torch.Tensor | None = None, + stiffness: torch.Tensor | None = None, + damping: torch.Tensor | None = None, + friction: torch.Tensor | None = None, + armature: torch.Tensor | None = None, + ) -> None: + """Read current parameters without tracking model gradients. + + See :meth:`ArticulationViewBase.fetch_joint_properties` for buffer + shapes, units and lifetime. + """ + # Configuration snapshots must not inherit gradients from Newton's + # differentiable model parameters. + with torch.no_grad(): + _checked_batch_call( + self.batch, + "fetch_joint_properties", + position_limits=position_limits, + velocity_limit=velocity_limit, + effort_limit=effort_limit, + stiffness=stiffness, + damping=damping, + friction=friction, + armature=armature, + ) + + def fetch_state( + self, + qpos: torch.Tensor, + qvel: torch.Tensor, + root_pose: torch.Tensor, + root_lin_vel: torch.Tensor, + root_ang_vel: torch.Tensor, + ) -> None: + """Read joint and root state through DexSim's existing batch operation. + + Buffer shapes and frames follow :meth:`ArticulationViewBase.fetch_state`. + Each call reads current simulation data, including writes since the + previous call; the reusable scratch buffer stores no cached state. + """ + scratch = self._root_pose_scratch + if scratch is None or scratch.shape != root_pose.shape: + scratch = torch.empty_like( + root_pose, dtype=torch.float32, device=self.device + ) + self._root_pose_scratch = scratch + _checked_batch_call( + self.batch, "fetch_state", qpos, qvel, scratch, root_lin_vel, root_ang_vel + ) + root_pose[..., :3].copy_(scratch[..., 4:]) + root_pose[..., 3:].copy_(scratch[..., :4]) + def fetch_root_pose(self, data: torch.Tensor) -> torch.Tensor: batch_pose = torch.empty_like(data, dtype=torch.float32, device=self.device) _checked_batch_call(self.batch, "fetch_root_pose", batch_pose) @@ -631,12 +689,10 @@ def apply_root_pose( rows = rows[changed] batch_pose = batch_pose[changed] - self._apply_rows( - "apply_root_pose", - batch_pose, - rows, - (7,), - ) + if len(rows): + _checked_batch_call( + self.batch, "apply_state", rows=rows, root_pose=batch_pose + ) def apply_root_velocity( self, velocity: torch.Tensor, env_ids: Sequence[int] | torch.Tensor @@ -649,9 +705,17 @@ def apply_root_velocity( Raises: RuntimeError: A batch write fails. Both velocity components are - restored first; a failed rollback is reported with the original - write error as its cause. + restored first on the default backend; a failed rollback is + reported with the original write error as its cause. Newton + uses the batch state writer, which validates both components + before writing and does not roll back execution failures. """ + if getattr(self.scene, "backend", None) == "newton": + # The batch state writer validates both components before mutation + # and propagates the final state once. No per-row batch or separate + # snapshot/linear/angular writes are needed for this backend. + self.apply_state(env_ids=env_ids, root_velocity=velocity) + return rows = self._select_rows(env_ids) expected_shape = (len(rows), 6) if tuple(velocity.shape) != expected_shape: @@ -696,53 +760,76 @@ def apply_root_velocity( def _joint_columns( self, joint_ids: Sequence[int] | torch.Tensor | None - ) -> torch.Tensor: + ) -> torch.Tensor | tuple[int, ...] | None: + # Static joint selections originate on the host. Keep them there for + # DexSim's topology validation instead of uploading then downloading + # the same indices on every control command. + if not isinstance(joint_ids, torch.Tensor): + columns = self._joint_dof_columns_host + if joint_ids is not None: + ids = tuple(int(index) for index in joint_ids) + if any(index < 0 or index >= len(columns) for index in ids): + raise IndexError(f"Joint selection is outside [0, {len(columns)}).") + columns = tuple(columns[index] for index in ids) + return None if columns == tuple(range(self.dof)) else columns return _selection_values( joint_ids, self._joint_dof_columns, label="Joint", ) - def _apply_joint_selection( + def apply_state( self, - values: torch.Tensor, - env_ids: Sequence[int] | torch.Tensor | None, - joint_ids: Sequence[int] | torch.Tensor | None, + env_ids: Sequence[int] | torch.Tensor | None = None, + joint_ids: Sequence[int] | torch.Tensor | None = None, *, - apply_method: str, - fetch_method: str, + root_pose: torch.Tensor | None = None, + qpos: torch.Tensor | None = None, + target_qpos: torch.Tensor | None = None, + qvel: torch.Tensor | None = None, + target_qvel: torch.Tensor | None = None, + qf: torch.Tensor | None = None, + root_velocity: torch.Tensor | None = None, + clear_dynamics: bool = False, ) -> None: - rows = self._select_rows(env_ids) - columns = self._joint_columns(joint_ids) - values = values.to(device=self.device, dtype=torch.float32) - expected = (len(rows), len(columns)) - if tuple(values.shape) != expected: - raise ValueError( - f"Expected selected joint data shape {expected}, got " - f"{tuple(values.shape)}." - ) - if not len(rows) or not len(columns): - return - if env_ids is None and joint_ids is None: - _checked_batch_call(self.batch, apply_method, values) - return + """Map EC fields once and submit one selected Scene batch write. - # DexSim's selected-articulation path currently materializes indices - # as host NumPy arrays. Keep the control path device-native by - # updating a reusable full-batch tensor and applying it without - # ``batch.select`` or ``dof_ids``. - scratch = self._joint_apply_scratch.get(apply_method) - scratch_shape = (self._row_count, self.batch.dof_width) - if scratch is None or tuple(scratch.shape) != scratch_shape: - scratch = torch.empty( - scratch_shape, - dtype=torch.float32, - device=self.device, + Shapes and selection semantics follow + :meth:`ArticulationViewBase.apply_state`. + """ + rows = None if env_ids is None else self._select_rows(env_ids) + count = self._row_count if rows is None else len(rows) + if count == 0: + return + fields = { + name: value + for name, value in ( + ("joint_position", qpos), + ("joint_target_position", target_qpos), + ("joint_velocity", qvel), + ("joint_target_velocity", target_qvel), + ("joint_force", qf), ) - self._joint_apply_scratch[apply_method] = scratch - _checked_batch_call(self.batch, fetch_method, scratch) - scratch[rows[:, None], columns] = values - _checked_batch_call(self.batch, apply_method, scratch) + if value is not None + } + columns = self._joint_columns(joint_ids) if fields else None + if root_pose is not None: + if tuple(root_pose.shape) != (count, 7): + raise ValueError(f"Expected root_pose shape {(count, 7)}.") + fields["root_pose"] = _batch_pose(root_pose.to(self.device, torch.float32)) + if root_velocity is not None: + if tuple(root_velocity.shape) != (count, 6): + raise ValueError(f"Expected root_velocity shape {(count, 6)}.") + fields["root_linear_velocity"] = root_velocity[:, :3] + fields["root_angular_velocity"] = root_velocity[:, 3:] + _checked_batch_call( + self.batch, + "apply_state", + rows=rows, + dof_ids=columns, + clear_dynamics=clear_dynamics, + **fields, + ) def apply_qpos( self, @@ -752,16 +839,8 @@ def apply_qpos( *, target: bool, ) -> None: - self._apply_joint_selection( - qpos, - env_ids, - joint_ids, - apply_method=( - "apply_joint_target_position" if target else "apply_joint_position" - ), - fetch_method=( - "fetch_joint_target_position" if target else "fetch_joint_position" - ), + self.apply_state( + env_ids, joint_ids, **{"target_qpos" if target else "qpos": qpos} ) def apply_qvel( @@ -772,16 +851,8 @@ def apply_qvel( *, target: bool, ) -> None: - self._apply_joint_selection( - qvel, - env_ids, - joint_ids, - apply_method=( - "apply_joint_target_velocity" if target else "apply_joint_velocity" - ), - fetch_method=( - "fetch_joint_target_velocity" if target else "fetch_joint_velocity" - ), + self.apply_state( + env_ids, joint_ids, **{"target_qvel" if target else "qvel": qvel} ) def apply_qf( @@ -790,19 +861,10 @@ def apply_qf( env_ids: Sequence[int] | torch.Tensor | None, joint_ids: Sequence[int] | torch.Tensor | None, ) -> None: - self._apply_joint_selection( - qf, - env_ids, - joint_ids, - apply_method="apply_joint_force", - fetch_method="fetch_joint_force", - ) + self.apply_state(env_ids, joint_ids, qf=qf) def clear_dynamics(self, env_ids: Sequence[int] | torch.Tensor) -> None: - rows = self._select_rows(env_ids) - if not len(rows): - return - _checked_batch_call(self.batch.select(rows), "clear_dynamics") + self.apply_state(env_ids, clear_dynamics=True) def compute_kinematics(self, env_ids: Sequence[int] | torch.Tensor) -> None: rows = self._select_rows(env_ids) diff --git a/embodichain/lab/sim/physics/newton.py b/embodichain/lab/sim/physics/newton.py index 8a684447a..cbb37d11b 100644 --- a/embodichain/lab/sim/physics/newton.py +++ b/embodichain/lab/sim/physics/newton.py @@ -102,6 +102,8 @@ def activate(self, sim_config: "SimulationManagerCfg") -> None: def sync_render_state(self, result: "dexsim.scene.Scene") -> None: """Publish Newton state through DexSim's render bridge without stepping.""" + if not self._manager.has_native_renderer: + return from dexsim.engine.newton_physics.backend_registry import get_newton_backend backend = get_newton_backend(result.world) diff --git a/embodichain/lab/sim/sensors/camera.py b/embodichain/lab/sim/sensors/camera.py index d2f5407a6..b7d9c5d41 100644 --- a/embodichain/lab/sim/sensors/camera.py +++ b/embodichain/lab/sim/sensors/camera.py @@ -146,6 +146,7 @@ def __init__( *, owner: SimulationManager, ) -> None: + owner._require_native_renderer("Native cameras") self._world = owner.get_world() self._arenas = [owner.get_env(i) for i in range(owner.num_envs)] if len(self._arenas) == 0: @@ -425,7 +426,9 @@ def look_at( env_ids (Sequence[int] | None): The environment IDs to set the look at for. If None, set for all environments. """ if up is None: - up = torch.tensor([[0.0, 0.0, 1.0]]).repeat(eye.shape[0], 1) + up = torch.tensor([[0.0, 0.0, 1.0]], device=eye.device).repeat( + eye.shape[0], 1 + ) pose = look_at_to_pose(eye, target, up) # To opengl coordinate system. @@ -482,28 +485,43 @@ def get_intrinsics(self) -> torch.Tensor: def reset(self, env_ids: Sequence[int] | None = None) -> None: self.cfg: CameraCfg + ids = ( + None + if env_ids is None + else torch.as_tensor(env_ids, dtype=torch.long, device=self.device) + ) if self.cfg.extrinsics.eye is not None: eye = ( - torch.tensor(self.cfg.extrinsics.eye, dtype=torch.float32) + torch.tensor( + self.cfg.extrinsics.eye, dtype=torch.float32, device=self.device + ) .squeeze_(0) .repeat(self.num_instances, 1) ) target = ( - torch.tensor(self.cfg.extrinsics.target, dtype=torch.float32) + torch.tensor( + self.cfg.extrinsics.target, dtype=torch.float32, device=self.device + ) .squeeze_(0) .repeat(self.num_instances, 1) ) up = ( - torch.tensor(self.cfg.extrinsics.up, dtype=torch.float32) + torch.tensor( + self.cfg.extrinsics.up, dtype=torch.float32, device=self.device + ) .squeeze_(0) .repeat(self.num_instances, 1) if self.cfg.extrinsics.up is not None else None ) + if ids is not None: + eye, target = eye[ids], target[ids] + if up is not None: + up = up[ids] self.look_at(eye, target, up, env_ids=env_ids) else: - pose = self.cfg.extrinsics.transformation + pose = self.cfg.extrinsics.transformation.to(self.device) pose = pose.unsqueeze_(0).repeat(self.num_instances, 1, 1) if self.cfg.extrinsics.parent is None: @@ -511,4 +529,4 @@ def reset(self, env_ids: Sequence[int] | None = None) -> None: pose[:, :3, 1] = -pose[:, :3, 1] pose[:, :3, 2] = -pose[:, :3, 2] - self.set_local_pose(pose, env_ids=env_ids) + self.set_local_pose(pose if ids is None else pose[ids], env_ids=env_ids) diff --git a/embodichain/lab/sim/sensors/contact_history.py b/embodichain/lab/sim/sensors/contact_history.py index 14fb4f893..944594d9b 100644 --- a/embodichain/lab/sim/sensors/contact_history.py +++ b/embodichain/lab/sim/sensors/contact_history.py @@ -36,6 +36,19 @@ __all__ = ["ContactHistory"] +def _current_warp_stream(device: torch.device) -> wp.Stream | None: + """Reuse a registered Warp stream when Torch is already running on it.""" + if device.type != "cuda": + return None + current = torch.cuda.current_stream(device) + stream = wp.get_stream(str(device)) + # A temporary second Warp wrapper unregisters the shared native handle on + # destruction in Warp 1.17, invalidating an enclosing capture. + if stream.cuda_stream == current.cuda_stream: + return stream + return wp.stream_from_torch(current) + + class ContactHistory: """Accumulate selected contacts across one control interval. @@ -132,11 +145,7 @@ def _launch(self, kernel, dim: int | tuple[int, ...], inputs: list) -> None: dim=dim, inputs=inputs, device=str(self.actor_ids.device), - stream=( - wp.stream_from_torch(self.actor_ids.device) - if self.actor_ids.is_cuda - else None - ), + stream=_current_warp_stream(self.actor_ids.device), ) def _finish_sample(self, dt: float) -> None: @@ -214,7 +223,13 @@ def reset(self, env_ids: Sequence[int] | torch.Tensor | None = None) -> None: Args: env_ids: Rows to clear. None selects every environment. """ - ids = slice(None) if env_ids is None else env_ids + ids = ( + None + if env_ids is None + else torch.as_tensor(env_ids, dtype=torch.long).to( + self.contact.device, non_blocking=True + ) + ) for value in ( self.contact, self.found, @@ -227,4 +242,7 @@ def reset(self, env_ids: Sequence[int] | torch.Tensor | None = None) -> None: self._hits, self._env_hits, ): - value[ids] = 0 + if ids is None: + value.zero_() + else: + value.index_fill_(0, ids, 0) diff --git a/embodichain/lab/sim/sensors/contact_sensor.py b/embodichain/lab/sim/sensors/contact_sensor.py index a08b5a121..74dfd3684 100644 --- a/embodichain/lab/sim/sensors/contact_sensor.py +++ b/embodichain/lab/sim/sensors/contact_sensor.py @@ -27,7 +27,10 @@ from embodichain.lab.sim.sensors import BaseSensor, SensorCfg from embodichain.lab.sim.sensors._warp.contact import scatter_contact_rows -from embodichain.lab.sim.sensors.contact_history import ContactHistory +from embodichain.lab.sim.sensors.contact_history import ( + ContactHistory, + _current_warp_stream, +) from embodichain.utils import configclass, logger if TYPE_CHECKING: @@ -161,6 +164,10 @@ def __init__( self._scatter_dropped_count ) self._histories: dict[str, ContactHistory] = {} + self._sample_graph = None + self._sample_graph_key = None + self._sample_stream = None + self._sample_torch_stream = None self._num_contacts_per_env: torch.Tensor | None = None """Number of contacts per environment.""" @@ -334,10 +341,12 @@ def update(self, **kwargs) -> None: self._scatter_dropped_count.copy_(self._sample_scatter_dropped_count) def _fetch_contacts(self) -> ContactBuffer: - assert self._query is not None and self._num_contacts_per_env is not None - self._num_contacts_per_env.zero_() - self._data_buffer["is_valid"].zero_() - self._sample_scatter_dropped_count.zero_() + buffer = self._read_contacts() + self._scatter_contacts(buffer) + return buffer + + def _read_contacts(self) -> ContactBuffer: + assert self._query is not None if self.device.type == "cuda": buffer = self._query.fetch_async() else: @@ -345,6 +354,13 @@ def _fetch_contacts(self) -> ContactBuffer: buffer.device_counts[0] = buffer.count buffer.device_counts[1] = buffer.dropped_count self._sync_filter_actor_metadata() + return buffer + + def _scatter_contacts(self, buffer: ContactBuffer) -> None: + assert self._num_contacts_per_env is not None + self._num_contacts_per_env.zero_() + self._data_buffer["is_valid"].zero_() + self._sample_scatter_dropped_count.zero_() wp.launch( kernel=scatter_contact_rows, dim=buffer.capacity, @@ -370,13 +386,8 @@ def _fetch_contacts(self) -> ContactBuffer: wp.from_torch(self._sample_scatter_dropped_count), ], device=str(self.device), - stream=( - wp.stream_from_torch(self.device) - if self.device.type == "cuda" - else None - ), + stream=_current_warp_stream(self.device), ) - return buffer @property def dropped_contacts(self) -> int: @@ -529,7 +540,54 @@ def update_physics_step(self, dt: float) -> None: Args: dt: Elapsed physics time in seconds. """ - buffer = self._fetch_contacts() + if dt <= 0: + raise ValueError("Contact sampling dt must be positive.") + buffer = self._read_contacts() + if self.device.type != "cuda" or torch.cuda.is_current_stream_capturing(): + self._sample_contacts(buffer, dt) + return + key = ( + id(buffer), + dt, + self.cfg.max_contacts_per_env, + tuple( + ( + id(h), + h.force_threshold, + h.include_unknown_counterpart, + h.counterpart_ids is None, + ) + for h in self._histories.values() + ), + ) + if self._sample_stream is None: + self._sample_stream = wp.get_stream(str(self.device)) + self._sample_torch_stream = wp.stream_to_torch(self._sample_stream) + caller = torch.cuda.current_stream(self.device) + same_stream = caller.cuda_stream == self._sample_torch_stream.cuda_stream + if not same_stream: + self._sample_torch_stream.wait_stream(caller) + with ( + wp.ScopedStream(self._sample_stream), + torch.cuda.stream(self._sample_torch_stream), + ): + if key != self._sample_graph_key: + self._sample_graph = None + self._sample_graph_key = key + # This real sample also compiles/warms every kernel. Do not + # execute history updates twice while preparing capture. + self._sample_contacts(buffer, dt) + else: + if self._sample_graph is None: + with wp.ScopedCapture(stream=self._sample_stream) as capture: + self._sample_contacts(buffer, dt) + self._sample_graph = capture.graph + wp.capture_launch(self._sample_graph, stream=self._sample_stream) + if not same_stream: + caller.wait_stream(self._sample_torch_stream) + + def _sample_contacts(self, buffer: ContactBuffer, dt: float) -> None: + self._scatter_contacts(buffer) self._query_dropped_count.add_(buffer.dropped_count_device) self._scatter_dropped_count.add_(self._sample_scatter_dropped_count) for history in self._histories.values(): @@ -541,15 +599,27 @@ def reset(self, env_ids: Sequence[int] | None = None) -> None: Args: env_ids: Rows to reset. None selects every row. """ - ids = slice(None) if env_ids is None else env_ids - self._data_buffer["is_valid"][ids] = False - self._num_contacts_per_env[ids] = 0 - self._scatter_dropped_count[ids] = 0 - self._sample_scatter_dropped_count[ids] = 0 + ids = ( + None + if env_ids is None + else torch.as_tensor(env_ids, dtype=torch.long).to( + self.device, non_blocking=True + ) + ) + for value in ( + self._data_buffer["is_valid"], + self._num_contacts_per_env, + self._scatter_dropped_count, + self._sample_scatter_dropped_count, + ): + if ids is None: + value.zero_() + else: + value.index_fill_(0, ids, 0) if env_ids is None: self._query_dropped_count.zero_() for history in self._histories.values(): - history.reset(env_ids) + history.reset(ids) def filter_by_user_ids( self, item_user_ids: torch.Tensor, env_ids: Sequence[int] | None = None diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index a180e0124..581188456 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -684,15 +684,17 @@ def __init__( # The structure is keys to the loaded texture data. The keys represent the texture group. self._texture_cache: Dict[str, Union[torch.Tensor, List[torch.Tensor]]] = dict() - self._init_sim_resources() - # The plane material and visibility are authored before declaration so # both eager Default loading and deferred Newton loading see them. self._spawn_default_plane_visibility = True + self._spawn_default_plane_material = None self._default_plane = None - self.set_default_background() + if self.has_native_renderer: + self._init_sim_resources() + self.set_default_background() self._declare_spawn_default_plane() - self.set_default_global_lighting() + if self.has_native_renderer: + self.set_default_global_lighting() # SpawnScene has already prepared the configured Arenas. Start the # optional browser runtime after default resources are declared. @@ -811,7 +813,7 @@ def set_default_renderer(cls, renderer: str = "auto", gpu_id: int = 0) -> str: from embodichain.lab.sim import cfg from embodichain.lab.sim.utility.render_utils import select_default_renderer - valid = {"auto", "hybrid", "fast-rt", "rt"} + valid = {"auto", "no-render", "hybrid", "fast-rt", "rt"} if renderer not in valid: logger.log_error( f"Invalid renderer '{renderer}'. Must be one of {sorted(valid)}." @@ -828,6 +830,17 @@ def set_default_renderer(cls, renderer: str = "auto", gpu_id: int = 0) -> str: logger.log_info(f"Default renderer set to '{resolved}'.") return resolved + @property + def has_native_renderer(self) -> bool: + """Whether the World owns a native renderer and visual resources.""" + return self.sim_config.render_cfg.renderer != "no-render" + + def _require_native_renderer(self, operation: str) -> None: + if not self.has_native_renderer: + raise RuntimeError( + f"{operation} requires a native renderer; the World uses no-render." + ) + @cached_property def num_envs(self) -> int: """Get the number of arenas in the simulation. @@ -1153,6 +1166,11 @@ def _convert_sim_config( sim_config.render_cfg.renderer = resolved_renderer sim_config.render_cfg.apply_to_dexsim_config(world_config) + if sim_config.render_cfg.renderer == "no-render": + if not sim_config.headless: + raise ValueError("renderer='no-render' requires headless=True.") + if getattr(sim_config.physics_cfg, "sync_to_renderer", None) is True: + raise ValueError("sync_to_renderer=True requires a native renderer.") if type(sim_config.device) is str: self.device = torch.device(sim_config.device) @@ -2068,7 +2086,8 @@ def _declare_spawn_default_plane(self) -> None: def _bind_default_plane(self, plane: Any) -> None: """Retain the spawned ground plane and apply its visibility.""" self._default_plane = plane - plane.set_visible(self._spawn_default_plane_visibility) + if self.has_native_renderer: + plane.set_visible(self._spawn_default_plane_visibility) def set_default_global_lighting(self) -> None: """Set default global lighting for the scene. @@ -2078,12 +2097,13 @@ def set_default_global_lighting(self) -> None: directional light is a global scene light (infinite distance) pointing downward along the -Z axis. """ + self._require_native_renderer("Default global lighting") # Environment emission light self.set_emission_light([1.0, 1.0, 1.0], 100.0) def set_default_background(self) -> None: """Set default background.""" - + self._require_native_renderer("Default background") mat_name = "plane_mat" mat_path = self._default_resources.get_material_path("PlaneDark") color_texture = os.path.join(mat_path, "PlaneDark_2K_Color.jpg") @@ -3188,7 +3208,7 @@ def get_robot(self, uid: str) -> Robot | None: def get_robot_uid_list(self) -> List[str]: """ - Retrieves a list of unique identifiers (UIDs) for all robots in the V2 system. + Retrieves a list of unique identifiers (UIDs) for all robots in the simulation system. Returns: list: A list containing the UIDs of the robots. @@ -4031,7 +4051,7 @@ def _capture_window_record_frame(self, state: _WindowRecordState) -> int: state.record_camera.render() rgb = np.asarray(state.record_camera.get_rgb_map()) if rgb.size != 0: - frame = np.ascontiguousarray(rgb[..., :3]) + frame = np.array(rgb[..., :3], copy=True, order="C") if frame is None: return state.task_status @@ -4478,6 +4498,7 @@ def create_visual_material(self, cfg: VisualMaterialCfg) -> VisualMaterial: VisualMaterial: the created visual material instance handle. """ + self._require_native_renderer("Native visual materials") if cfg.uid in self._visual_materials: logger.log_warning( f"Visual material {cfg.uid} already exists. Returning the existing one." diff --git a/embodichain/learning/rl/algo/ppo.py b/embodichain/learning/rl/algo/ppo.py index d94c4c8a1..9f56daab5 100644 --- a/embodichain/learning/rl/algo/ppo.py +++ b/embodichain/learning/rl/algo/ppo.py @@ -83,7 +83,9 @@ def __init__(self, cfg: PPOCfg, policy): def update(self, rollout: TensorDict) -> Dict[str, float]: """Update the policy using a collected rollout.""" - rollout = rollout.clone() + # GAE adds fields but does not mutate the collected tensors. A shallow + # clone isolates those fields without copying the entire GPU rollout. + rollout = rollout.clone(recurse=False) compute_gae(rollout, gamma=self.cfg.gamma, gae_lambda=self.cfg.gae_lambda) flat_rollout = transition_view(rollout, flatten=True) diff --git a/embodichain/learning/rl/policy_evaluation/cli.py b/embodichain/learning/rl/policy_evaluation/cli.py index 120c6eb38..83e1ca5fd 100644 --- a/embodichain/learning/rl/policy_evaluation/cli.py +++ b/embodichain/learning/rl/policy_evaluation/cli.py @@ -105,7 +105,7 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: parser.add_argument("--physics-backend") parser.add_argument( "--renderer", - choices=("hybrid", "fast-rt", "rt"), + choices=("no-render", "hybrid", "fast-rt", "rt"), ) parser.add_argument("--gpu-id", type=int, default=0) parser.add_argument("--scene-config") @@ -361,6 +361,10 @@ def _build_native_runtime( if viewer and not uses_simulator: raise ValueError("--viewer requires a simulator training task") renderer = args.renderer or str(trainer.get("renderer", "hybrid")) + if viewer and renderer == "no-render": + if args.renderer is not None: + raise ValueError("--viewer requires a native renderer.") + renderer = "hybrid" num_envs = ( 1 if viewer @@ -406,6 +410,8 @@ def _build_native_runtime( def _validate_native_options(args: argparse.Namespace) -> None: + if args.viewer and args.renderer == "no-render": + raise ValueError("--viewer requires a native renderer.") profile_options = { "--resource-root": args.resource_root, "--physics-backend": args.physics_backend, diff --git a/embodichain/learning/rl/utils/trainer.py b/embodichain/learning/rl/utils/trainer.py index 8ac4124f1..942650647 100644 --- a/embodichain/learning/rl/utils/trainer.py +++ b/embodichain/learning/rl/utils/trainer.py @@ -95,6 +95,7 @@ def __init__( self.device = self.algorithm.device self.global_step = 0 + self._training_start_step = 0 self.num_updates = 0 self.start_time = time.time() self.ret_window = deque(maxlen=100) @@ -178,6 +179,8 @@ def _pack_log_dict(self, prefix: str, data: dict) -> dict: return out def train(self, total_timesteps: int) -> dict[str, Any]: + self._training_start_step = self.global_step + self.start_time = time.time() if self.rank == 0: print(f"Start training, total steps: {total_timesteps}") num_envs = int(self.env.num_envs) @@ -238,11 +241,12 @@ def on_step(tensordict: TensorDict, info: dict): metrics_dict = info.get("metrics") self._log_scalar_dict("rewards", rewards_dict) self._log_scalar_dict("metrics", metrics_dict) - log_dict = {} - log_dict.update(self._pack_log_dict("rewards", rewards_dict)) - log_dict.update(self._pack_log_dict("metrics", metrics_dict)) - if log_dict and self.use_wandb: - wandb.log(log_dict, step=self.global_step) + if self.use_wandb: + log_dict = {} + log_dict.update(self._pack_log_dict("rewards", rewards_dict)) + log_dict.update(self._pack_log_dict("metrics", metrics_dict)) + if log_dict: + wandb.log(log_dict, step=self.global_step) rollout = self.buffer.start_rollout() rollout = self.collector.collect( @@ -326,7 +330,7 @@ def _sync_episode_stats(self) -> None: def _log_train(self, losses: dict[str, float]): elapsed = max(1e-6, time.time() - self.start_time) - sps = self.global_step / elapsed + sps = (self.global_step - self._training_start_step) / elapsed avgR = np.mean(self.ret_window) if len(self.ret_window) > 0 else float("nan") avgL = np.mean(self.len_window) if len(self.len_window) > 0 else float("nan") history_entry = { @@ -463,12 +467,15 @@ def save_checkpoint(self, path: str | None = None) -> str | None: return path def get_summary(self) -> dict[str, Any]: + """Return persistent counters and throughput for the current train call.""" elapsed = max(1e-6, time.time() - self.start_time) return { "global_step": int(self.global_step), "num_updates": int(self.num_updates), "elapsed_time_sec": float(elapsed), - "training_fps": float(self.global_step / elapsed), + "training_fps": float( + (self.global_step - self._training_start_step) / elapsed + ), "last_train_metrics": dict(self.last_train_metrics), "last_eval_metrics": dict(self.last_eval_metrics), "train_history": list(self.train_history), diff --git a/embodichain/utils/math.py b/embodichain/utils/math.py index 0df8bf467..1c77bc49b 100644 --- a/embodichain/utils/math.py +++ b/embodichain/utils/math.py @@ -26,35 +26,27 @@ def look_at_to_pose( - eye: Union[torch.Tensor, list], - target: Union[torch.Tensor, list], - up: Union[torch.Tensor, list] = [0, 0, 1], + eye: torch.Tensor | list, + target: torch.Tensor | list, + up: torch.Tensor | list = [0, 0, 1], ) -> torch.Tensor: - """Get the camera pose from eye to target with up direction, supporting batch processing. + """Return batched camera poses on the eye tensor's device. Args: - eye (Union[torch.Tensor, list]): Camera positions with shape (N, 3). - target (Union[torch.Tensor, list]): Target positions with shape (N, 3). - up (Union[torch.Tensor, list], optional): Up directions with shape (N, 3) or (3,). Defaults to [0, 0, 1]. + eye: Camera positions with shape (N, 3). + target: Target positions with shape (N, 3). + up: Up directions with shape (N, 3) or (3,). Defaults to [0, 0, 1]. Returns: - torch.Tensor: Camera pose matrices with shape (N, 4, 4). + Camera pose matrices with shape (N, 4, 4), on the eye tensor's device. """ eye = ( torch.tensor(eye, dtype=torch.float32) if not isinstance(eye, torch.Tensor) else eye ) - target = ( - torch.tensor(target, dtype=torch.float32) - if not isinstance(target, torch.Tensor) - else target - ) - up = ( - torch.tensor(up, dtype=torch.float32) - if not isinstance(up, torch.Tensor) - else up - ) + target = torch.as_tensor(target, dtype=eye.dtype, device=eye.device) + up = torch.as_tensor(up, dtype=eye.dtype, device=eye.device) if eye.ndim == 1: eye = eye.unsqueeze(0) @@ -80,7 +72,7 @@ def look_at_to_pose( camera_x_norm = torch.norm(camera_x, dim=1, keepdim=True) if torch.any(camera_x_norm < 1e-6): # Handle degenerate cases up = ( - torch.tensor([0, 1, 0], dtype=torch.float32) + torch.tensor([0, 1, 0], dtype=eye.dtype, device=eye.device) .unsqueeze(0) .repeat(eye.shape[0], 1) ) @@ -92,7 +84,9 @@ def look_at_to_pose( # Construct camera pose matrices camera_pose = ( - torch.eye(4, dtype=torch.float32).unsqueeze(0).repeat(eye.shape[0], 1, 1) + torch.eye(4, dtype=torch.float32, device=eye.device) + .unsqueeze(0) + .repeat(eye.shape[0], 1, 1) ) # (N, 4, 4) camera_pose[:, :3, 0] = camera_x camera_pose[:, :3, 1] = camera_y diff --git a/embodichain_tasks/configs/tasks/locomotion/velocity/g1_flat/agents/ppo.newton.yaml b/embodichain_tasks/configs/tasks/locomotion/velocity/g1_flat/agents/ppo.newton.yaml index 28af09ead..e8538f685 100644 --- a/embodichain_tasks/configs/tasks/locomotion/velocity/g1_flat/agents/ppo.newton.yaml +++ b/embodichain_tasks/configs/tasks/locomotion/velocity/g1_flat/agents/ppo.newton.yaml @@ -4,7 +4,7 @@ trainer: seed: 42 device: cuda:0 headless: true - renderer: hybrid + renderer: no-render num_envs: 4096 iterations: 10001 buffer_size: 24 diff --git a/embodichain_tasks/configs/tasks/locomotion/velocity/g1_flat/agents/ppo.yaml b/embodichain_tasks/configs/tasks/locomotion/velocity/g1_flat/agents/ppo.yaml index d1c3f36cc..6e18d95aa 100644 --- a/embodichain_tasks/configs/tasks/locomotion/velocity/g1_flat/agents/ppo.yaml +++ b/embodichain_tasks/configs/tasks/locomotion/velocity/g1_flat/agents/ppo.yaml @@ -4,7 +4,7 @@ trainer: seed: 42 device: cuda:0 headless: true - renderer: hybrid + renderer: no-render num_envs: 4096 iterations: 10001 buffer_size: 24 diff --git a/embodichain_tasks/configs/tasks/locomotion/velocity/g1_flat/env.newton.yaml b/embodichain_tasks/configs/tasks/locomotion/velocity/g1_flat/env.newton.yaml index 474a1006f..8d1393c7f 100644 --- a/embodichain_tasks/configs/tasks/locomotion/velocity/g1_flat/env.newton.yaml +++ b/embodichain_tasks/configs/tasks/locomotion/velocity/g1_flat/env.newton.yaml @@ -2,12 +2,11 @@ physics: newton id: UnitreeG1FlatRL-v1 max_episode_steps: 1000 num_envs: 8 -scene_node_capacity: 262144 headless: true device: cuda:0 arena_space: 2.5 render_cfg: - renderer: hybrid + renderer: no-render physics_config: sync_to_renderer: false gravity: @@ -19,7 +18,7 @@ physics_config: solver_type: mujoco_warp njmax: 1024 nconmax: 256 - iterations: 10 + iterations: 20 ls_iterations: 50 integrator: implicitfast separate_worlds: true diff --git a/embodichain_tasks/configs/tasks/locomotion/velocity/g1_flat/env.yaml b/embodichain_tasks/configs/tasks/locomotion/velocity/g1_flat/env.yaml index 537aa0ad9..bb0c86376 100644 --- a/embodichain_tasks/configs/tasks/locomotion/velocity/g1_flat/env.yaml +++ b/embodichain_tasks/configs/tasks/locomotion/velocity/g1_flat/env.yaml @@ -2,12 +2,11 @@ physics: default id: UnitreeG1FlatRL-v1 max_episode_steps: 1000 num_envs: 8 -scene_node_capacity: 262144 headless: true device: cuda:0 arena_space: 2.5 render_cfg: - renderer: hybrid + renderer: no-render physics_config: cache_material: true gravity: diff --git a/embodichain_tasks/embodichain_tasks/locomotion/_robot.py b/embodichain_tasks/embodichain_tasks/locomotion/_robot.py index 31836437a..1ef580c01 100644 --- a/embodichain_tasks/embodichain_tasks/locomotion/_robot.py +++ b/embodichain_tasks/embodichain_tasks/locomotion/_robot.py @@ -59,6 +59,7 @@ def resolve_or_default(name: str, field_name: str): return None return JointDrivePropertiesCfg( drive_type="force", + target_mode=None if defaults is None else defaults.target_mode, stiffness=stiffness, damping=damping, max_effort=resolve_or_default("effort_limit", "max_effort"), diff --git a/embodichain_tasks/embodichain_tasks/locomotion/velocity/_embodichain.py b/embodichain_tasks/embodichain_tasks/locomotion/velocity/_embodichain.py index af6b0da3c..c11c2f9d7 100644 --- a/embodichain_tasks/embodichain_tasks/locomotion/velocity/_embodichain.py +++ b/embodichain_tasks/embodichain_tasks/locomotion/velocity/_embodichain.py @@ -28,9 +28,11 @@ from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg from embodichain.learning.rl.policy_evaluation.camera import PolicyViewerCameraCfg +from embodichain.lab.sim.types import EnvObs from embodichain.utils.math import quat_apply, quat_apply_inverse from .._robot import apply_task_joint_drive_properties +from .contracts._math import constant __all__ = ["EmbodiChainVelocityEnv"] @@ -80,6 +82,20 @@ class EmbodiChainVelocityEnv(EmbodiedEnv): explicit_pd_effort_control = False def __init__(self, cfg: EmbodiedEnvCfg | None = None, **kwargs) -> None: + rewards = ( + self.velocity_task_config.data.get("rewards", {}) + if self.velocity_task_config is not None + else {} + ) + self._state_fields = frozenset( + field.name + for field in fields(self.state_type) + if not ( + field.name == "self_collision_count" + and field.default is None + and "self_collisions" not in rewards + ) + ) cfg = EmbodiedEnvCfg() if cfg is None else cfg if cfg.robot is not MISSING and self.velocity_task_config is not None: term_cfg = ( @@ -168,7 +184,6 @@ def _init_sim_state(self, **kwargs) -> None: device=self.device, ) sensor = self.get_sensor("locomotion_contacts") - robot_ids = sensor.get_actor_ids(self.robot.uid, self.robot.link_names) self._foot_contacts = sensor.create_history( "feet", sensor.get_actor_ids(self.robot.uid, self.foot_link_names), @@ -181,17 +196,19 @@ def _init_sim_state(self, **kwargs) -> None: "illegal", sensor.get_actor_ids(self.robot.uid, self.illegal_contact_link_names), ) - self._self_contacts = sensor.create_history( - "self", - robot_ids, - counterpart_ids=robot_ids, - force_threshold=float( - config.data.get("rewards", {}) - .get("self_collisions", {}) - .get("params", {}) - .get("force_threshold", 0.0) - ), - ) + if "self_collision_count" in self._state_fields: + robot_ids = sensor.get_actor_ids(self.robot.uid, self.robot.link_names) + self._self_contacts = sensor.create_history( + "self", + robot_ids, + counterpart_ids=robot_ids, + force_threshold=float( + config.data.get("rewards", {}) + .get("self_collisions", {}) + .get("params", {}) + .get("force_threshold", 0.0) + ), + ) self._locomotion_generator = self.get_generator("locomotion.commands_and_noise") action_shape = (self.num_envs, config.action_dim) foot_shape = (self.num_envs, len(self.foot_link_names)) @@ -208,6 +225,7 @@ def _init_sim_state(self, **kwargs) -> None: self._command_steps_remaining = torch.zeros( self.num_envs, dtype=torch.long, device=self.device ) + self._command_steps_until_check = 0 self._bind_locomotion_action_state() bias = config.data.get("events", {}).get("encoder_bias") if bias is not None: @@ -346,6 +364,7 @@ def set_velocity_command( self.command.copy_(torch.as_tensor(values, device=self.device).unsqueeze(0)) self.reward_command.copy_(self.command) self._command_steps_remaining.fill_(torch.iinfo(torch.long).max) + self._command_steps_until_check = torch.iinfo(torch.long).max self._heading_env.fill_(False) self._forward_env.fill_(False) static_standing = stopped and (frequency is None or frequency == 0.0) @@ -388,7 +407,15 @@ def _active_command_ranges(self) -> torch.Tensor: tuple(stage.get(name, ranges[name])) for name in ("lin_vel_x", "lin_vel_y", "ang_vel_z") ) - return torch.tensor(active, dtype=torch.float32, device=self.device) + key = (active, self.device) + cached = getattr(self, "_command_ranges_cache", None) + if cached is None or cached[0] != key: + cached = ( + key, + torch.tensor(active, dtype=torch.float32, device=self.device), + ) + self._command_ranges_cache = cached + return cached[1] def _active_single_axis_fraction(self) -> float: command = self.velocity_task_config.data["commands"]["twist"] @@ -440,6 +467,12 @@ def _resample_commands(self, env_ids: torch.Tensor) -> None: self._command_steps_remaining[env_ids] = ( torch.ceil(seconds / self.step_dt).to(torch.long).clamp_min(1) ) + # Partial reset may shorten the earliest deadline. The one-step slack + # keeps the host bound conservative under float32 interval rounding. + self._command_steps_until_check = min( + self._command_steps_until_check, + max(1, math.floor(float(minimum) / self.step_dt) - 1), + ) heading_range = command["ranges"].get("heading") if bool(command.get("heading_command")) and heading_range is not None: self._heading_target[env_ids] = torch.empty( @@ -453,31 +486,49 @@ def _resample_commands(self, env_ids: torch.Tensor) -> None: count, device=self.device, generator=self._locomotion_generator ) <= float(command.get("rel_heading_envs", 0.0)) else: - self._heading_env[env_ids] = False + self._heading_env.index_fill_(0, env_ids, False) self._standing_env[env_ids] = torch.rand( count, device=self.device, generator=self._locomotion_generator ) <= float(command.get("rel_standing_envs", 0.0)) gait_frequency_range = command.get("gait_frequency_range") if gait_frequency_range is None: - self._gait_frequency[env_ids] = 0.0 + self._gait_frequency.index_fill_(0, env_ids, 0.0) else: - self._gait_frequency[env_ids] = torch.empty( - count, device=self.device - ).uniform_( + frequency = torch.empty(count, device=self.device).uniform_( float(gait_frequency_range[0]), float(gait_frequency_range[1]), generator=self._locomotion_generator, ) - self._gait_frequency[env_ids[self._standing_env[env_ids]]] = 0.0 + frequency.masked_fill_(self._standing_env[env_ids], 0.0) + self._gait_frequency[env_ids] = frequency self._forward_env[env_ids] = torch.rand( count, device=self.device, generator=self._locomotion_generator ) <= float(command.get("rel_forward_envs", 0.0)) - forward_ids = env_ids[self._forward_env[env_ids]] - if forward_ids.numel() > 0: - self.command[forward_ids, 0] = ( - self.command[forward_ids, 0].abs().clamp_min(0.3) + forward = self._forward_env[env_ids] + selected_command = self.command[env_ids] + selected_command[:, 0].copy_( + torch.where( + forward, + selected_command[:, 0].abs().clamp_min(0.3), + selected_command[:, 0], ) - self.command[forward_ids, 1:] = 0.0 + ) + selected_command[:, 1:].masked_fill_(forward.unsqueeze(-1), 0.0) + self.command[env_ids] = selected_command + + def _advance_command_timer(self) -> None: + """Advance GPU timers without selecting rows before any can be due.""" + self._command_steps_remaining.sub_(1) + self._command_steps_until_check -= 1 + if self._command_steps_until_check > 0: + return + resample_ids = ( + (self._command_steps_remaining <= 0).nonzero(as_tuple=False).squeeze(-1) + ) + self._resample_commands(resample_ids) + self._command_steps_until_check = int( + self._command_steps_remaining.min().item() + ) def _update_heading_commands(self) -> None: command = self.velocity_task_config.data["commands"]["twist"] @@ -491,43 +542,54 @@ def _update_heading_commands(self) -> None: error = torch.atan2(torch.sin(error), torch.cos(error)) yaw_range = self._active_command_ranges()[2] yaw = (float(command.get("heading_control_stiffness", 0.0)) * error).clamp( - min=float(yaw_range[0]), max=float(yaw_range[1]) + min=yaw_range[0], max=yaw_range[1] + ) + self.command[:, 2].copy_( + torch.where(self._heading_env, yaw, self.command[:, 2]) ) - self.command[self._heading_env, 2] = yaw[self._heading_env] - self.command[self._standing_env] = 0.0 + self.command.masked_fill_(self._standing_env.unsqueeze(-1), 0.0) def _update_sim_state(self, **kwargs) -> None: super()._update_sim_state(**kwargs) - velocity = self.robot.get_qvel()[:, self.policy_joint_ids] - self._joint_acceleration.copy_( - (velocity - self._previous_joint_velocity) / self.step_dt - ) - self._previous_joint_velocity.copy_(velocity) - root_quaternion = self.robot.body_data.root_pose[:, 3:7] - root_lin_vel_w = self.robot.body_data.root_lin_vel - root_ang_vel_w = self.robot.body_data.root_ang_vel - root_velocity_w = torch.cat((root_lin_vel_w, root_ang_vel_w), dim=-1) - self._root_acceleration_w.copy_( - (root_velocity_w - self._previous_root_velocity_w) / self.step_dt - ) - self._previous_root_velocity_w.copy_(root_velocity_w) - base_lin_vel_b = quat_apply_inverse(root_quaternion, root_lin_vel_w) - base_ang_vel_b = quat_apply_inverse(root_quaternion, root_ang_vel_w) - filter_weight = float( - self.velocity_task_config.data.get("velocity_filter_weight", 1.0) - ) - self._filtered_base_lin_vel_b.lerp_(base_lin_vel_b, filter_weight) - self._filtered_base_ang_vel_b.lerp_(base_ang_vel_b, filter_weight) - self._update_foot_swing_height() + state = None + if { + "joint_acc", + "root_acceleration_w", + "reward_base_lin_vel_b", + "reward_base_ang_vel_b", + } & self._state_fields: + state = self.robot.body_data.fetch_state() + if "joint_acc" in self._state_fields: + velocity = state["qvel"][:, self.policy_joint_ids] + self._joint_acceleration.copy_( + (velocity - self._previous_joint_velocity) / self.step_dt + ) + self._previous_joint_velocity.copy_(velocity) + if "root_acceleration_w" in self._state_fields: + root_velocity_w = torch.cat( + (state["root_lin_vel"], state["root_ang_vel"]), + dim=-1, + ) + self._root_acceleration_w.copy_( + (root_velocity_w - self._previous_root_velocity_w) / self.step_dt + ) + self._previous_root_velocity_w.copy_(root_velocity_w) + if {"reward_base_lin_vel_b", "reward_base_ang_vel_b"} & self._state_fields: + root_quaternion = state["root_pose"][:, 3:7] + base_lin_vel_b = quat_apply_inverse(root_quaternion, state["root_lin_vel"]) + base_ang_vel_b = quat_apply_inverse(root_quaternion, state["root_ang_vel"]) + filter_weight = float( + self.velocity_task_config.data.get("velocity_filter_weight", 1.0) + ) + self._filtered_base_lin_vel_b.lerp_(base_lin_vel_b, filter_weight) + self._filtered_base_ang_vel_b.lerp_(base_ang_vel_b, filter_weight) + if "foot_swing_height_cost" in self._state_fields: + self._update_foot_swing_height() self._episode_step += 1 self._global_control_step += 1 self._gait_process.add_(self.step_dt * self._gait_frequency).remainder_(1.0) self.reward_command.copy_(self.command) - self._command_steps_remaining.sub_(1) - resample_ids = ( - (self._command_steps_remaining <= 0).nonzero(as_tuple=False).squeeze(-1) - ) - self._resample_commands(resample_ids) + self._advance_command_timer() self._update_heading_commands() self._invalidate_task_cache() @@ -568,34 +630,26 @@ def _initialize_episode( ) pose[:, 6] = torch.cos(0.5 * yaw) pose[:, 5] = torch.sin(0.5 * yaw) - self.robot.clear_dynamics(env_ids=ids) - self.robot.set_local_pose(pose, env_ids=ids) - joint_position = torch.as_tensor( - self.velocity_task_config.default_joint_position, - dtype=torch.float32, - device=self.device, + joint_position = constant( + self.velocity_task_config.default_joint_position, pose ).expand(count, -1) - self.robot.set_qpos( - joint_position, - joint_ids=self.policy_joint_ids, + self.robot.set_state( env_ids=ids, - target=False, - ) - self.robot.set_qpos( - joint_position, joint_ids=self.policy_joint_ids, - env_ids=ids, - target=True, - ) - self._episode_step[ids] = 0 - self._previous_joint_velocity[ids] = 0.0 - self._joint_acceleration[ids] = 0.0 - self._previous_root_velocity_w[ids] = 0.0 - self._root_acceleration_w[ids] = 0.0 - self._filtered_base_lin_vel_b[ids] = 0.0 - self._filtered_base_ang_vel_b[ids] = 0.0 - self._foot_peak_height[ids] = 0.0 - self._foot_swing_height_cost[ids] = 0.0 + root_pose=pose, + qpos=joint_position, + target_qpos=joint_position, + clear_dynamics=True, + ) + self._episode_step.index_fill_(0, ids, 0) + self._previous_joint_velocity.index_fill_(0, ids, 0.0) + self._joint_acceleration.index_fill_(0, ids, 0.0) + self._previous_root_velocity_w.index_fill_(0, ids, 0.0) + self._root_acceleration_w.index_fill_(0, ids, 0.0) + self._filtered_base_lin_vel_b.index_fill_(0, ids, 0.0) + self._filtered_base_ang_vel_b.index_fill_(0, ids, 0.0) + self._foot_peak_height.index_fill_(0, ids, 0.0) + self._foot_swing_height_cost.index_fill_(0, ids, 0.0) self._resample_commands(ids) self._update_heading_commands() self.reward_command[ids] = self.command[ids] @@ -655,10 +709,11 @@ def _update_foot_swing_height(self) -> None: ) def _common_state(self) -> dict[str, torch.Tensor]: - root_pose = self.robot.body_data.root_pose + state = self.robot.body_data.fetch_state() + root_pose = state["root_pose"] root_quaternion = root_pose[:, 3:7] - root_lin_vel_w = self.robot.body_data.root_lin_vel - root_ang_vel_w = self.robot.body_data.root_ang_vel + root_lin_vel_w = state["root_lin_vel"] + root_ang_vel_w = state["root_ang_vel"] gravity_w = torch.zeros_like(root_lin_vel_w) gravity_w[:, 2] = -1.0 link_pose = self.robot.body_data.body_link_pose @@ -693,7 +748,7 @@ def _common_state(self) -> dict[str, torch.Tensor]: orientation_velocity, self._imu_offset, ) - return { + common = { "base_lin_vel_b": base_lin_vel_b, "base_ang_vel_b": base_ang_vel_b, "reward_base_lin_vel_b": self._filtered_base_lin_vel_b, @@ -703,10 +758,9 @@ def _common_state(self) -> dict[str, torch.Tensor]: "reward_command": self.reward_command, "gait_frequency": self._gait_frequency, "gait_process": self._gait_process, - "joint_pos": self.robot.get_qpos()[:, self.policy_joint_ids], - "joint_vel": self.robot.get_qvel()[:, self.policy_joint_ids], + "joint_pos": state["qpos"][:, self.policy_joint_ids], + "joint_vel": state["qvel"][:, self.policy_joint_ids], "joint_acc": self._joint_acceleration, - "joint_torque": self.robot.get_qf()[:, self.policy_joint_ids], "action": self.locomotion_action, "last_action": self.last_locomotion_action, "episode_step": self._episode_step, @@ -725,9 +779,6 @@ def _common_state(self) -> dict[str, torch.Tensor]: "soft_joint_lower": (midpoint - half_range).unsqueeze(0), "soft_joint_upper": (midpoint + half_range).unsqueeze(0), "reward_body_ang_vel_w": orientation_velocity[:, 3:], - "angular_momentum_w": self._orbital_angular_momentum( - link_pose[:, :, :3], link_velocity[:, :, :3] - ), "illegal_contact_force": ( illegal_force_by_body.amax(dim=-1) if illegal_force_by_body.shape[-1] @@ -737,11 +788,20 @@ def _common_state(self) -> dict[str, torch.Tensor]: "orientation_projected_gravity_b": quat_apply_inverse( orientation_pose[:, 3:7], gravity_w ), - "self_collision_count": self._self_contacts.contact_count, "encoder_bias": self.encoder_bias, "foot_swing_height_cost": self._foot_swing_height_cost, } + if "joint_torque" in self._state_fields: + common["joint_torque"] = self.robot.get_qf()[:, self.policy_joint_ids] + if "angular_momentum_w" in self._state_fields: + common["angular_momentum_w"] = self._orbital_angular_momentum( + link_pose[:, :, :3], link_velocity[:, :, :3] + ) + if "self_collision_count" in self._state_fields: + common["self_collision_count"] = self._self_contacts.contact_count + return common + def _make_task_state(self, common: dict[str, torch.Tensor]) -> Any: values: dict[str, Any] = {} for field in fields(self.state_type): @@ -759,17 +819,43 @@ def get_velocity_locomotion_state(self) -> Any: self._state_cache = self._make_task_state(self._common_state()) return self._state_cache + def _extend_obs(self, obs: EnvObs, **kwargs) -> EnvObs: + # Actor and critic functors consume the same task snapshot. Build the + # clean pair once within this manager pass, then discard it so resets, + # later observations and direct callers cannot reuse an old result. + previous = getattr(self, "_observation_pair_cache", None) + self._observation_pair_cache = [] + try: + return super()._extend_obs(obs, **kwargs) + finally: + self._observation_pair_cache = previous + def build_velocity_locomotion_observations( self, *, enable_corruption: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: - """Build actor and asymmetric critic observations.""" - actor, critic = self.build_observations_fn( - self.velocity_task_config, self.get_velocity_locomotion_state() - ) - if enable_corruption and self.corrupt_actor_fn is not None: + """Build independent actor and critic tensors, with fresh requested noise. + + Observation-manager terms share one clean pair per pass. Calls outside + that pass build a new pair from the current task state. + """ + cache = getattr(self, "_observation_pair_cache", None) + if cache is None or not cache: + pair = self.build_observations_fn( + self.velocity_task_config, self.get_velocity_locomotion_state() + ) + if cache is not None: + cache.append(pair) + else: + pair = cache[0] + actor, critic = pair + corrupt = enable_corruption and self.corrupt_actor_fn is not None + if cache is not None: + actor, critic = actor.clone(), critic.clone() + elif corrupt: actor = actor.clone() + if corrupt: actor = self.corrupt_actor_fn(actor, self._locomotion_generator) return actor, critic diff --git a/embodichain_tasks/embodichain_tasks/locomotion/velocity/contracts/_math.py b/embodichain_tasks/embodichain_tasks/locomotion/velocity/contracts/_math.py index abba3a841..fa6e83105 100644 --- a/embodichain_tasks/embodichain_tasks/locomotion/velocity/contracts/_math.py +++ b/embodichain_tasks/embodichain_tasks/locomotion/velocity/contracts/_math.py @@ -35,7 +35,11 @@ def constant(values: tuple[float, ...], like: torch.Tensor) -> torch.Tensor: """Create a constant tensor on the same device and dtype as ``like``.""" - return torch.tensor(values, device=like.device, dtype=like.dtype) + # The direct CUDA constructor waits for the upload, which also stalls on + # earlier work in the stream. Queue the copy on the caller's stream while + # preserving independent storage for each invocation. + host = torch.tensor(values, device="cpu", dtype=like.dtype) + return host.to(device=like.device, non_blocking=True) def phase_signal( diff --git a/embodichain_tasks/embodichain_tasks/locomotion/velocity/contracts/_reward_terms.py b/embodichain_tasks/embodichain_tasks/locomotion/velocity/contracts/_reward_terms.py index f7238d0b1..e772fd766 100644 --- a/embodichain_tasks/embodichain_tasks/locomotion/velocity/contracts/_reward_terms.py +++ b/embodichain_tasks/embodichain_tasks/locomotion/velocity/contracts/_reward_terms.py @@ -154,7 +154,7 @@ def reward_terms( pose_error = torch.square(state.joint_pos - default) gait_cfg = rewards["foot_gait"]["params"] - offsets = state.joint_pos.new_tensor(gait_cfg["offset"]) + offsets = constant(tuple(gait_cfg["offset"]), state.joint_pos) global_phase = ( state.episode_step * control_dt / float(gait_cfg["period"]) ).unsqueeze(-1) diff --git a/tests/gym/envs/tasks/test_locomotion_deployments.py b/tests/gym/envs/tasks/test_locomotion_deployments.py index f7dcab223..57765ce84 100644 --- a/tests/gym/envs/tasks/test_locomotion_deployments.py +++ b/tests/gym/envs/tasks/test_locomotion_deployments.py @@ -64,7 +64,13 @@ def test_packaged_deployment_and_agent_config_agree( sim_cfg.DefaultPhysicsCfg if backend == "default" else sim_cfg.NewtonPhysicsCfg ) assert isinstance(cfg.sim_cfg.physics_cfg, physics_type) - assert cfg.sim_cfg.render_cfg.renderer in {"auto", "hybrid", "fast-rt", "rt"} + assert cfg.sim_cfg.render_cfg.renderer in { + "auto", + "hybrid", + "fast-rt", + "rt", + "no-render", + } assert agent["trainer"]["renderer"] == cfg.sim_cfg.render_cfg.renderer assert Path(agent["trainer"]["gym_config"]).resolve() == path.resolve() assert set(agent) == {"trainer", "policy", "algorithm"} diff --git a/tests/gym/envs/tasks/test_unitree_locomotion.py b/tests/gym/envs/tasks/test_unitree_locomotion.py index 96df5f268..0a6c2478b 100644 --- a/tests/gym/envs/tasks/test_unitree_locomotion.py +++ b/tests/gym/envs/tasks/test_unitree_locomotion.py @@ -193,7 +193,8 @@ def test_unitree_deployments_preserve_task_physics( ) assert isinstance(config.sim_cfg.physics_cfg, physics_type) assert config.sim_cfg.physics_dt == pytest.approx(0.005) - assert config.sim_cfg.scene_node_capacity == 262144 + expected_capacity = None if robot == "g1" else 262144 + assert config.sim_cfg.scene_node_capacity == expected_capacity if backend == "default": assert config.sim_cfg.physics_cfg.to_dexsim_args()["cache_material"] is True else: diff --git a/tests/sim/objects/test_articulation_state.py b/tests/sim/objects/test_articulation_state.py new file mode 100644 index 000000000..62fb4f827 --- /dev/null +++ b/tests/sim/objects/test_articulation_state.py @@ -0,0 +1,65 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from embodichain.lab.sim.objects.articulation import Articulation + +pytestmark = pytest.mark.no_sim + + +def test_combined_state_clamps_both_positions_before_one_write() -> None: + art = object.__new__(Articulation) + art.device = torch.device("cpu") + art._entities = [object()] * 3 + art.__dict__["dof"] = 2 + art._data = SimpleNamespace( + qpos_limits=torch.tensor([[[-1.0, 1.0], [-2.0, 2.0]]] * 3), + articulation_view=SimpleNamespace(apply_state=Mock()), + ) + art._stabilize_newton_mimic_target_write = Mock() + art.set_state( + env_ids=[2, 0], + joint_ids=[1, 0], + qpos=torch.tensor([[5.0, -3.0], [-5.0, 3.0]]), + target_qpos=torch.tensor([[-4.0, 4.0], [4.0, -4.0]]), + clear_dynamics=True, + ) + call = art._data.articulation_view.apply_state + call.assert_called_once() + rows, columns = call.call_args.args + assert torch.as_tensor(rows).tolist() == [2, 0] + assert torch.as_tensor(columns).tolist() == [1, 0] + torch.testing.assert_close( + call.call_args.kwargs["qpos"], torch.tensor([[2.0, -1.0], [-2.0, 1.0]]) + ) + torch.testing.assert_close( + call.call_args.kwargs["target_qpos"], torch.tensor([[-2.0, 1.0], [2.0, -1.0]]) + ) + call.reset_mock() + with pytest.raises(ValueError, match="root_velocity"): + art.set_state( + qpos=torch.zeros((3, 2)), + root_velocity=torch.zeros((3, 5)), + clear_dynamics=True, + ) + call.assert_not_called() diff --git a/tests/sim/objects/test_scene_backend.py b/tests/sim/objects/test_scene_backend.py index 7095ef2dd..fb53f5474 100644 --- a/tests/sim/objects/test_scene_backend.py +++ b/tests/sim/objects/test_scene_backend.py @@ -19,6 +19,7 @@ import importlib.util import inspect from types import SimpleNamespace +from unittest.mock import Mock import pytest import torch @@ -109,6 +110,9 @@ def test_scene_views_match_installed_dexsim_batch_surface() -> None: "apply_joint_force", "fetch_joint_force", "clear_dynamics", + "fetch_joint_properties", + "fetch_state", + "apply_state", "fetch_joint_acceleration", "fetch_link_pose", "fetch_link_linear_velocity", @@ -327,6 +331,96 @@ def __init__(self) -> None: self.root_pose_apply_rows: list[tuple[int, ...]] = [] self.clear_dynamics_rows: list[tuple[int, ...]] = [] self.selections: list[tuple[int, ...]] = [] + self.joint_properties = { + "position_limits": torch.tensor( + [ + [[-1.0, 1.0], [-2.0, 2.0], [-3.0, 3.0]], + [[-4.0, 4.0], [-5.0, 5.0], [-6.0, 6.0]], + ] + ), + **{ + field: self.position + offset + for offset, field in enumerate( + ( + "velocity_limit", + "effort_limit", + "stiffness", + "damping", + "friction", + "armature", + ) + ) + }, + } + self.property_reads: list[tuple[str, ...]] = [] + self.property_status = len(self) + + def fetch_joint_properties(self, **outputs: torch.Tensor | None) -> int: + selected = {name: out for name, out in outputs.items() if out is not None} + self.property_reads.append(tuple(selected)) + for name, out in selected.items(): + out.copy_(self.joint_properties[name]) + return self.property_status + + def fetch_state( + self, + qpos: torch.Tensor, + qvel: torch.Tensor, + pose: torch.Tensor, + linear: torch.Tensor, + angular: torch.Tensor, + ) -> int: + for out, source in zip( + (qpos, qvel, pose, linear, angular), + ( + self.position, + self.velocity, + self.root_pose, + self.root_linear_velocity, + self.root_angular_velocity, + ), + strict=True, + ): + out.copy_(source) + return len(self) + + def apply_state( + self, + *, + rows: torch.Tensor | None = None, + dof_ids: torch.Tensor | None = None, + clear_dynamics: bool = False, + **fields: torch.Tensor, + ) -> int: + rows = ( + torch.arange(len(self)) if rows is None else rows.as_subclass(torch.Tensor) + ) + columns = ( + torch.arange(self.dof_width) + if dof_ids is None + else dof_ids.as_subclass(torch.Tensor) + ) + self.last_dof_ids = tuple(columns.tolist()) + selected = _SelectedArticulationBatch(self, rows) + if "joint_position" in fields: + self.position[rows[:, None], columns] = fields["joint_position"] + if clear_dynamics: + selected.clear_dynamics() + joint_fields = { + "joint_position": "position", + "joint_target_position": "target_position", + "joint_velocity": "velocity", + "joint_target_velocity": "target_velocity", + "joint_force": "force", + } + for name, value in fields.items(): + if name == "joint_position": + continue + elif name in joint_fields: + getattr(self, joint_fields[name])[rows[:, None], columns] = value + else: + getattr(selected, f"apply_{name}")(value) + return len(rows) def __len__(self) -> int: return len(self.force) @@ -461,6 +555,109 @@ def test_articulation_data_uses_scene_view(backend: str) -> None: assert scene.articulation_batch_objects == entities +@pytest.mark.parametrize("backend", ["dexsim", "newton"]) +def test_articulation_properties_use_batch_and_preserve_per_env_values( + backend: str, +) -> None: + scene = _Scene(backend) + scalar = Mock(side_effect=AssertionError("per-articulation property read")) + entities = [ + SimpleNamespace( + get_joint_position_limits=scalar, + get_joint_velocity_limit=scalar, + get_joint_effort_limit=scalar, + get_drive=scalar, + get_newton_drive=scalar, + ) + for _ in range(2) + ] + data = ArticulationData(entities, scene, torch.device("cpu")) + batch = scene.articulation_batch + assert batch.property_reads == [ + ("position_limits", "velocity_limit", "effort_limit") + ] + for value, name in ( + (data.qpos_limits, "position_limits"), + (data.qvel_limits, "velocity_limit"), + (data.qf_limits, "effort_limit"), + ): + torch.testing.assert_close(value, batch.joint_properties[name]) + for name in ("stiffness", "damping", "friction", "armature"): + previous = getattr(data, f"joint_{name}") + expected = batch.joint_properties[name].clone() + batch.joint_properties[name].add_(10.0) + current = getattr(data, f"joint_{name}") + torch.testing.assert_close(previous, expected) + torch.testing.assert_close(current, batch.joint_properties[name]) + scalar.assert_not_called() + + +def test_articulation_property_reads_do_not_track_model_parameter_gradients() -> None: + scene = _Scene("newton") + for value in scene.articulation_batch.joint_properties.values(): + value.requires_grad_() + + data = ArticulationData([_ArticulationEntity()] * 2, scene, torch.device("cpu")) + for output, name in ( + (data.qpos_limits, "position_limits"), + (data.qvel_limits, "velocity_limit"), + (data.qf_limits, "effort_limit"), + *( + (getattr(data, f"joint_{name}"), name) + for name in ("stiffness", "damping", "friction", "armature") + ), + ): + source = scene.articulation_batch.joint_properties[name] + torch.testing.assert_close(output, source) + assert not output.requires_grad + assert output.grad_fn is None + assert source.requires_grad + + +def test_articulation_property_initialization_reports_batch_failure() -> None: + scene = _Scene() + scene.articulation_batch.property_status = -1 + with pytest.raises(RuntimeError, match="fetch_joint_properties"): + ArticulationData([_ArticulationEntity()] * 2, scene, torch.device("cpu")) + + +@pytest.mark.parametrize("backend", ["dexsim", "newton"]) +def test_articulation_state_read_preserves_layout_and_refreshes_after_write( + backend: str, +) -> None: + scene = _Scene(backend) + batch = scene.articulation_batch + batch.fetch_state = Mock(wraps=batch.fetch_state) + data = ArticulationData( + [_ArticulationEntity(), _ArticulationEntity()], scene, torch.device("cpu") + ) + first = data.fetch_state() + for key, source in ( + ("root_pose", _embodichain_pose(batch.root_pose)), + ("qpos", batch.position), + ("qvel", batch.velocity), + ("root_lin_vel", batch.root_linear_velocity), + ("root_ang_vel", batch.root_angular_velocity), + ): + torch.testing.assert_close(first[key], source) + unchanged = {key: value[0].clone() for key, value in first.items()} + batch.apply_state( + rows=torch.tensor([1]), + root_pose=torch.tensor([[0.0, 0.0, 0.0, 1.0, 7.0, 8.0, 9.0]]), + joint_position=torch.tensor([[4.0, 5.0, 6.0]]), + clear_dynamics=True, + ) + second = data.fetch_state() + for key, value in second.items(): + torch.testing.assert_close(value[0], unchanged[key]) + torch.testing.assert_close(second["root_pose"], _embodichain_pose(batch.root_pose)) + torch.testing.assert_close(second["qpos"], batch.position) + torch.testing.assert_close(second["qvel"][1], torch.zeros(3)) + reads = batch.fetch_state.call_args_list + assert len(reads) == 2 + assert reads[0].args[2].data_ptr() == reads[1].args[2].data_ptr() + + def test_newton_articulation_geometry_merges_every_render_mesh() -> None: scene = _Scene("newton") entities = [_NewtonArticulationEntity(), _NewtonArticulationEntity()] @@ -694,7 +891,7 @@ def test_articulation_partial_force_preserves_other_rows_and_dofs() -> None: torch.tensor([[1.0, 2.0, 3.0], [4.0, 50.0, 6.0]]), ) assert batch.selections == [] - assert batch.last_dof_ids is None + assert batch.last_dof_ids == (1,) def test_articulation_joint_mapping_stays_on_the_view_device() -> None: diff --git a/tests/sim/sensors/test_camera_reset.py b/tests/sim/sensors/test_camera_reset.py new file mode 100644 index 000000000..deef1cea9 --- /dev/null +++ b/tests/sim/sensors/test_camera_reset.py @@ -0,0 +1,78 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from collections.abc import Sequence + +import numpy as np +import pytest +import torch + +from embodichain.lab.sim.sensors.camera import Camera, CameraCfg + +pytestmark = pytest.mark.no_sim + + +class _View: + def __init__(self) -> None: + self.pose = np.zeros((4, 4), dtype=np.float32) + + def set_local_pose(self, pose: np.ndarray) -> None: + self.pose = pose.copy() + + +@pytest.mark.parametrize("extrinsics_kind", ["pose", "look_at", "look_at_default_up"]) +@pytest.mark.parametrize( + "env_ids", + [None, [1], (1,), (2, 0), (), [], torch.tensor([2, 0])], + ids=[ + "all", + "list", + "tuple_one", + "tuple_ordered", + "tuple_empty", + "list_empty", + "tensor", + ], +) +def test_reset_preserves_batch_dimension_and_unselected_views( + extrinsics_kind: str, env_ids: Sequence[int] | torch.Tensor | None +) -> None: + """Sequence indices select camera rows for both extrinsics representations.""" + if extrinsics_kind == "pose": + extrinsics = CameraCfg.ExtrinsicsCfg(pos=(1.0, 2.0, 3.0)) + else: + extrinsics = CameraCfg.ExtrinsicsCfg( + eye=(1.0, 2.0, 3.0), + target=(0.0, 0.0, 0.0), + up=(0.0, 0.0, 1.0) if extrinsics_kind == "look_at" else None, + ) + camera = Camera.__new__(Camera) + camera.cfg = CameraCfg(uid="camera", extrinsics=extrinsics) + camera.device = torch.device("cpu") + camera._entities = [_View() for _ in range(3)] + camera._num_instances = len(camera._entities) + camera.reset(env_ids) + + expected = camera.cfg.extrinsics.transformation.clone() + expected[:3, 1:3].neg_() + selected = set(range(3) if env_ids is None else map(int, env_ids)) + for row, view in enumerate(camera._entities): + torch.testing.assert_close( + torch.from_numpy(view.pose), + expected if row in selected else torch.zeros((4, 4)), + ) diff --git a/tests/sim/sensors/test_contact_sensor_graph.py b/tests/sim/sensors/test_contact_sensor_graph.py new file mode 100644 index 000000000..466377c71 --- /dev/null +++ b/tests/sim/sensors/test_contact_sensor_graph.py @@ -0,0 +1,161 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +import warp as wp + +from dexsim.scene import ContactActorInfo, ContactBuffer, ContactQueryCapabilities +from embodichain.lab.sim.sensors import ContactSensor, ContactSensorCfg + +pytestmark = [pytest.mark.gpu, pytest.mark.no_sim] + + +class _CudaQuery: + def __init__(self, device: torch.device) -> None: + self.capabilities = ContactQueryCapabilities(True, True, True) + self.selected_actor_ids = (10, 20) + self.actors = ( + ContactActorInfo(10, "arena_0/cube", None, "arena_0", 0), + ContactActorInfo(20, "arena_1/cube", None, "arena_1", 1), + ) + self.buffer = ContactBuffer.allocate(4, str(device)) + self.buffer.data.zero_() + self.buffer.data[:2, 5] = 1.0 + self.buffer.data[:2, 9] = 0.2 + self.buffer.actor_ids[:2] = torch.tensor( + [[0, 10], [0, 20]], device=device, dtype=torch.int32 + ) + self.buffer.env_ids[:2] = torch.tensor([0, 1], device=device, dtype=torch.int32) + self.buffer.device_counts.zero_() + + def actor_info(self, actor_id: int) -> ContactActorInfo: + return next(actor for actor in self.actors if actor.actor_id == actor_id) + + def fetch_async(self) -> ContactBuffer: + return self.buffer + + +@pytest.mark.parametrize("caller_kind", ["warp", "torch"]) +def test_sensor_graph_samples_once_reuses_and_recaptures(caller_kind: str) -> None: + """Capture and replay retain timing, reset isolation and stream ordering.""" + wp.init() + device = torch.device("cuda:0") + caller = ( + wp.stream_to_torch(wp.get_stream(str(device))) + if caller_kind == "warp" + else torch.cuda.Stream(device=device) + ) + with torch.cuda.stream(caller): + query = _CudaQuery(device) + owner = SimpleNamespace( + num_envs=2, + arena_offsets=torch.zeros((2, 3), device=device), + _spawn_scene=SimpleNamespace( + handles=lambda uid: tuple( + SimpleNamespace(path=actor.path) for actor in query.actors + ) + ), + spawn_result=SimpleNamespace( + create_contact_query=lambda *args, **kwargs: query + ), + ) + sensor = ContactSensor( + ContactSensorCfg( + uid="contacts", rigid_uid_list=["cube"], max_contacts_per_env=2 + ), + device, + owner=owner, + ) + history = sensor.create_history( + "feet", torch.tensor([[10], [20]], device=device) + ) + sensor.begin_control_step() + + def assert_air_time(value: float) -> None: + torch.testing.assert_close( + history.current_air_time, + torch.full((2, 1), value, device=device), + ) + + sensor.update_physics_step(0.1) + assert sensor._sample_graph is None + assert_air_time(0.1) + sensor.update_physics_step(0.1) + graph = sensor._sample_graph + assert graph is not None + assert_air_time(0.2) + sensor.update_physics_step(0.1) + assert sensor._sample_graph is graph + assert_air_time(0.3) + + query.buffer.device_counts[0] = 2 + sensor.update_physics_step(0.1) + assert sensor._sample_graph is graph + assert history.contact.all() + assert history.found.all() + assert history.first_contact.all() + assert_air_time(0.0) + torch.testing.assert_close( + history.last_air_time, torch.full((2, 1), 0.4, device=device) + ) + torch.testing.assert_close( + history.force, + torch.tensor([[[0.0, 0.0, 2.0]], [[0.0, 0.0, 2.0]]], device=device), + ) + torch.testing.assert_close(history.contact_count, torch.ones(2, device=device)) + assert sensor.get_data()["is_valid"].tolist() == [[True, False], [True, False]] + fields = ( + "contact", + "found", + "first_contact", + "force", + "peak_force", + "current_air_time", + "last_air_time", + "contact_count", + ) + untouched = {name: getattr(history, name)[1].clone() for name in fields} + valid = sensor.get_data()["is_valid"][1].clone() + sensor.reset((0,)) + for name in fields: + assert not getattr(history, name)[0].any() + torch.testing.assert_close(getattr(history, name)[1], untouched[name]) + assert not sensor.get_data()["is_valid"][0].any() + torch.testing.assert_close(sensor.get_data()["is_valid"][1], valid) + assert sensor._num_contacts_per_env.tolist() == [0, 1] + assert sensor._sample_graph is graph + + sensor.begin_control_step() + history.force_threshold = 3.0 + sensor.update_physics_step(0.1) + assert sensor._sample_graph is None + assert not history.contact.any() + assert not history.found.any() + assert_air_time(0.1) + sensor.update_physics_step(0.1) + recaptured = sensor._sample_graph + assert recaptured is not None and recaptured is not graph + assert_air_time(0.2) + sensor.update_physics_step(0.1) + assert sensor._sample_graph is recaptured + assert not history.contact.any() + assert_air_time(0.3) + caller.synchronize() diff --git a/tests/sim/test_cfg.py b/tests/sim/test_cfg.py index 47ac1bb64..e7ad85b02 100644 --- a/tests/sim/test_cfg.py +++ b/tests/sim/test_cfg.py @@ -594,9 +594,7 @@ def names(config_type: type) -> set[str]: "collision_filter_parent", "is_visible", "is_site", - # DexSim 0.5.1rc1 native loader/ordering controls remain outside the - # backend-neutral EmbodiChain collision configuration. - "priority", + # Asset-loading adapter option, not a native ShapeConfig parameter. "use_native_mesh_loader", } assert ( diff --git a/tests/sim/test_no_renderer.py b/tests/sim/test_no_renderer.py new file mode 100644 index 000000000..83158f861 --- /dev/null +++ b/tests/sim/test_no_renderer.py @@ -0,0 +1,154 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Fresh-process checks for the physics-only simulation configuration.""" + +from __future__ import annotations + +import os +from pathlib import Path +import subprocess +import sys + +import pytest + + +@pytest.mark.subprocess_sim +@pytest.mark.gpu +@pytest.mark.parametrize("backend", ["default", "newton"]) +def test_no_render_world_steps_and_resets(backend: str, tmp_path: Path) -> None: + urdf = tmp_path / "slider.urdf" + urdf.write_text( + """ + + + + + + + + + + + + + + + + + + """, + encoding="utf-8", + ) + result = subprocess.run( + [sys.executable, str(Path(__file__).resolve()), backend, str(urdf)], + env=os.environ.copy(), + text=True, + capture_output=True, + timeout=90, + ) + assert result.returncode == 0, result.stdout + result.stderr + + +def _run_world(backend: str, urdf_path: str) -> None: + import torch + from embodichain.lab.sim import SimulationManager, SimulationManagerCfg + from embodichain.lab.sim.cfg import ( + ArticulationCfg, + DefaultPhysicsCfg, + NewtonPhysicsCfg, + RenderCfg, + RigidObjectCfg, + ) + from embodichain.lab.sim.shapes import CubeCfg + + physics = ( + NewtonPhysicsCfg( + device="cuda:0", + num_substeps=1, + physics_dt=0.005, + solver_cfg={ + "solver_type": "mujoco_warp", + "iterations": 20, + "ls_iterations": 50, + }, + ) + if backend == "newton" + else DefaultPhysicsCfg(device="cpu", physics_dt=0.005) + ) + cfg = SimulationManagerCfg( + headless=True, + num_envs=2, + physics_cfg=physics, + render_cfg=RenderCfg(renderer="no-render"), + startup_summary="off", + ) + sim = SimulationManager(cfg) + body = articulation = None + try: + body = sim.add_rigid_object( + RigidObjectCfg( + uid="box", + shape=CubeCfg(), + init_pos=(0.0, 0.0, 1.0), + ) + ) + sim.prepare() + initial = body.get_local_pose(to_matrix=True).clone() + for _ in range(10): + sim.update(step=1) + after = body.get_local_pose(to_matrix=True).clone() + assert torch.isfinite(after).all() + assert torch.all(after[:, 2, 3] < initial[:, 2, 3]) + body.reset(env_ids=[0]) + reset_pose = body.get_local_pose(to_matrix=True) + torch.testing.assert_close(reset_pose[0], initial[0]) + torch.testing.assert_close(reset_pose[1], after[1]) + # The collision ground remains present, even though it has no material. + assert sim._spawn_scene.handles("default_plane") + assert not sim._world.is_window_initialized() + assert not sim.has_native_renderer + assert not sim._visual_materials + from embodichain.lab.sim.sensors import CameraCfg + + with pytest.raises(RuntimeError, match="requires a native renderer"): + sim.add_sensor(CameraCfg(uid="camera", width=32, height=32)) + assert not sim._sensors + + # Exercise the URDF source resolver against the selected DexSim build. + articulation = sim.add_articulation( + ArticulationCfg(uid="slider", fpath=urdf_path, init_pos=(0.0, 0.0, 2.0)) + ) + sim.prepare() + assert articulation.dof == 1 + state = torch.full((2, 1), 0.05, device=sim.device) + articulation.set_qpos(state, target=False) + torch.testing.assert_close(articulation.body_data.qpos, state) + articulation.reset(env_ids=[0]) + torch.testing.assert_close( + articulation.body_data.qpos[0], torch.zeros_like(state[0]) + ) + torch.testing.assert_close(articulation.body_data.qpos[1], state[1]) + sim.update(step=1) + assert torch.isfinite(articulation.body_data.qpos).all() + finally: + body = articulation = None + sim.destroy(exit_process=False) + SimulationManager.flush_cleanup_queue() + + +if __name__ == "__main__": + _run_world(sys.argv[1], sys.argv[2]) diff --git a/tests/sim/test_sim_manager.py b/tests/sim/test_sim_manager.py index d47680d68..833cf81bb 100644 --- a/tests/sim/test_sim_manager.py +++ b/tests/sim/test_sim_manager.py @@ -1944,7 +1944,9 @@ def test_add_camera_uses_owning_manager_render_context( arenas = [object(), object()] sim = object.__new__(SimulationManager) - sim.sim_config = SimpleNamespace(num_envs=len(arenas)) + sim.sim_config = SimpleNamespace( + num_envs=len(arenas), render_cfg=RenderCfg(renderer="hybrid") + ) sim.device = torch.device("cpu") sim._world = world sim._arenas = arenas diff --git a/tests/sim/test_sim_manager_cfg.py b/tests/sim/test_sim_manager_cfg.py index 47abd7f7f..6adf8c294 100644 --- a/tests/sim/test_sim_manager_cfg.py +++ b/tests/sim/test_sim_manager_cfg.py @@ -375,20 +375,33 @@ def test_newton_teardown_skips_cpu_devices( render_sync.clear.assert_called_once_with() +@pytest.mark.no_sim +@pytest.mark.parametrize("has_native_renderer", [True, False]) def test_newton_backend_uses_unified_render_sync_entry( monkeypatch: pytest.MonkeyPatch, + has_native_renderer: bool, ) -> None: world = object() native_backend = SimpleNamespace(sync_to_dexsim=MagicMock()) + lookup = MagicMock( + side_effect=lambda candidate: native_backend if candidate is world else None + ) monkeypatch.setattr( "dexsim.engine.newton_physics.backend_registry.get_newton_backend", - lambda candidate: native_backend if candidate is world else None, + lookup, + ) + backend = NewtonPhysicsBackend( + SimpleNamespace(has_native_renderer=has_native_renderer) ) - backend = NewtonPhysicsBackend(SimpleNamespace()) backend.sync_render_state(SimpleNamespace(world=world)) - native_backend.sync_to_dexsim.assert_called_once_with(world) + if has_native_renderer: + lookup.assert_called_once_with(world) + native_backend.sync_to_dexsim.assert_called_once_with(world) + else: + lookup.assert_not_called() + native_backend.sync_to_dexsim.assert_not_called() @pytest.mark.parametrize( diff --git a/tests/utils/test_math.py b/tests/utils/test_math.py index 47b985843..426d8f3c0 100644 --- a/tests/utils/test_math.py +++ b/tests/utils/test_math.py @@ -24,6 +24,7 @@ convert_quat, default_orientation, inv_transform, + look_at_to_pose, matrix_from_quat, quat_apply, quat_conjugate, @@ -36,6 +37,24 @@ ) +@pytest.mark.parametrize( + "device", ["cpu", pytest.param("cuda:0", marks=pytest.mark.gpu)] +) +def test_look_at_keeps_vertical_camera_pose_on_eye_device(device: str) -> None: + eye = torch.tensor([[0.0, 0.0, 1.0], [1.0, 0.0, 1.0]], device=device) + + pose = look_at_to_pose(eye, [[0.0, 0.0, 0.0]] * 2) + + assert pose.device == eye.device + assert torch.isfinite(pose).all() + torch.testing.assert_close(pose[:, :3, 3], eye) + rotation = pose[:, :3, :3] + torch.testing.assert_close( + rotation.transpose(1, 2) @ rotation, + torch.eye(3, device=device).expand(2, 3, 3), + ) + + @pytest.mark.parametrize("backend", ["numpy", "torch"]) @pytest.mark.parametrize("dtype", [np.float32, np.float64]) @pytest.mark.parametrize("kind", ["identity", "translation", "rotation_translation"])