From 4164eae27ad2f701e1edb75fca3d69db783f99a6 Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Mon, 26 Oct 2026 03:22:17 +0800 Subject: [PATCH 1/4] feat: integrate tensor ray observations --- pyproject.rocm.toml | 10 + pyproject.toml | 11 + src/unilab/base/entity.py | 504 +++++++++++++++--- src/unilab/envs/manager_based_rl_env.py | 31 ++ src/unilab/envs/mdp/__init__.py | 4 + src/unilab/envs/mdp/ray_observation.py | 197 +++++++ src/unilab/managers/observation_manager.py | 71 ++- src/unilab/managers/types.py | 3 + tests/base/test_entity_tensor_read_plan.py | 159 +++++- tests/envs/mdp/test_ray_observation.py | 278 ++++++++++ tests/managers/test_device_contracts.py | 13 + tests/scripts/test_rocm_dependency_profile.py | 20 + 12 files changed, 1211 insertions(+), 90 deletions(-) create mode 100644 src/unilab/envs/mdp/ray_observation.py create mode 100644 tests/envs/mdp/test_ray_observation.py diff --git a/pyproject.rocm.toml b/pyproject.rocm.toml index 6ff905e47..f10e8c29a 100644 --- a/pyproject.rocm.toml +++ b/pyproject.rocm.toml @@ -73,6 +73,10 @@ mujoco = [ "mjbatch-uni~=0.2.5", ] motrix = ["unisim-core[motrix]==2.1.1"] +# Branch-local ROCm integration against sibling ray-contract checkouts. Keep +# this extra aligned with pyproject.toml; the default runtime still resolves the +# published unisim-core pin. uni-ray never becomes a runtime dependency. +ray-local = ["uni-ray"] [dependency-groups] # The RL stack (APPO/off-policy/multi-GPU) is merged into unilab itself @@ -91,6 +95,12 @@ torch = [ triton-rocm = [ { index = "pytorch-rocm72", marker = "sys_platform == 'linux' and platform_machine == 'x86_64'" }, ] +unisim-core = [ + { path = "../unisim", editable = true, marker = "extra == 'ray-local'" }, +] +uni-ray = [ + { path = "../uni_ray", editable = true, marker = "extra == 'ray-local'" }, +] [tool.uv] exclude-dependencies = [ diff --git a/pyproject.toml b/pyproject.toml index c2e991376..5a5440364 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -113,6 +113,11 @@ genesis = ["genesis-world==1.3.3"] superdex = [ "unisim-core[superdex]==2.1.1 ; python_version >= '3.12' and sys_platform == 'linux' and platform_machine == 'x86_64'", ] +# Branch-local integration against sibling ray-contract checkouts. The exact +# published UniSim pin remains the default; the local path source overrides it +# only when this optional extra is selected. uni-ray never becomes a runtime +# dependency and is imported only by UniSim's public caster factory. +ray-local = ["uni-ray"] [dependency-groups] # The RL stack (APPO/off-policy/multi-GPU) is merged into unilab itself @@ -129,6 +134,12 @@ explicit = true torch = [ { index = "r2-cu130", marker = "sys_platform=='linux' or sys_platform=='win32'" }, ] +unisim-core = [ + { path = "../unisim", editable = true, marker = "extra == 'ray-local'" }, +] +uni-ray = [ + { path = "../uni_ray", editable = true, marker = "extra == 'ray-local'" }, +] [tool.uv] exclude-dependencies = ["torchvision"] required-environments = [ diff --git a/src/unilab/base/entity.py b/src/unilab/base/entity.py index 854ca79f9..54c163243 100644 --- a/src/unilab/base/entity.py +++ b/src/unilab/base/entity.py @@ -7,6 +7,7 @@ from __future__ import annotations +import inspect import re from collections.abc import Callable, Iterator, Mapping, Sequence from contextlib import contextmanager @@ -205,11 +206,26 @@ def __post_init__(self) -> None: class EntityTensorBodyStateView: """Entity-ordered world-frame Torch body state for one read phase.""" + body_names: tuple[str, ...] + pos_w: torch.Tensor | None + quat_w: torch.Tensor | None + lin_vel_w: torch.Tensor | None + ang_vel_w: torch.Tensor | None + + +@dataclass(frozen=True) +class SceneTensorPoseView: + """Ray-shaped full-batch Torch pose backed by packed scene tensor views. + + This is deliberately ``RayPoseTensorView``-shaped without importing the + caster package. The packet-backed tensors may be row-strided column views; + the public caster accepts such packed-strided Torch views without requiring + global contiguity. + """ + body_names: tuple[str, ...] pos_w: torch.Tensor quat_w: torch.Tensor - lin_vel_w: torch.Tensor - ang_vel_w: torch.Tensor _TENSOR_BODY_SENSOR_FIELDS: tuple[tuple[str, str], ...] = ( @@ -218,23 +234,51 @@ class EntityTensorBodyStateView: ("lin_vel_w", "track_linvel_w_"), ("ang_vel_w", "track_angvel_w_"), ) +_TENSOR_POSE_BODY_FIELDS = ("pos_w", "quat_w") +_TENSOR_BODY_FIELDS = frozenset((*_TENSOR_POSE_BODY_FIELDS, "lin_vel_w", "ang_vel_w")) +_TENSOR_BODY_FIELD_ORDER = ("pos_w", "quat_w", "lin_vel_w", "ang_vel_w") +_TENSOR_BODY_PACKET_FIELDS = { + "pos_w": "body_pos_w", + "quat_w": "body_quat_w", + "lin_vel_w": "body_lin_vel_w", + "ang_vel_w": "body_ang_vel_w", +} +_TENSOR_PACKET_FIELD_WIDTHS = { + "body_pos_w": 3, + "body_quat_w": 4, + "body_lin_vel_w": 3, + "body_ang_vel_w": 3, +} +_PACKET_FIELD_BODY_FIELDS = { + packet_name: field for field, packet_name in _TENSOR_BODY_PACKET_FIELDS.items() +} _TRACKED_BODY_SENSOR_PREFIXES = {field: prefix for field, prefix in _TENSOR_BODY_SENSOR_FIELDS} _TENSOR_SENSOR_WIDTHS = {"pos_w": 3, "quat_w": 4, "lin_vel_w": 3, "ang_vel_w": 3} +def _packet_body_name(entity_name: str, body_name: str) -> str: + return f"{entity_name}/{body_name}" + + @dataclass(frozen=True) class SceneTensorReadSpec: """Declare one entity's named state reads for the scene tensor phase. ``sensor_names`` are explicit backend names. ``body_names`` are entity-local - names; the scene owner expands them into the canonical ``track_*`` sensor - contract. Requests are explicit so the packed packet never grows from an - implicit EntityCfg namespace that is only needed by a cold-path consumer. + names; ``body_fields`` selects their public semantic blocks. The legacy + default is the exact four-field body-state set; pose-only ray consumers + select ``(pos_w, quat_w)`` so velocities never enter the packed transfer. """ entity: str sensor_names: tuple[str, ...] = () body_names: tuple[str, ...] = () + body_fields: tuple[str, ...] = ( + "pos_w", + "quat_w", + "lin_vel_w", + "ang_vel_w", + ) def __post_init__(self) -> None: if not isinstance(self.entity, str) or not self.entity: @@ -243,8 +287,11 @@ def __post_init__(self) -> None: raise TypeError("Scene tensor read sensor_names must be a sequence of strings") if isinstance(self.body_names, (str, bytes)): raise TypeError("Scene tensor read body_names must be a sequence of strings") + if isinstance(self.body_fields, (str, bytes)): + raise TypeError("Scene tensor read body_fields must be a sequence of strings") sensors = tuple(self.sensor_names) bodies = tuple(self.body_names) + fields = tuple(self.body_fields) if not sensors and not bodies: # Generic joint observations request the entity's packed qpos/qvel # packet without named sensors or body-projected state. This is a @@ -255,8 +302,19 @@ def __post_init__(self) -> None: raise TypeError(f"Scene tensor read {label} names must be non-empty strings") if len(set(names)) != len(names): raise ValueError(f"Scene tensor read {label} names must be unique: {names}") + if bodies and not fields: + raise ValueError("Scene tensor read body_fields cannot be empty when bodies are set") + unknown = [field for field in fields if field not in _TENSOR_BODY_FIELDS] + if unknown: + raise ValueError( + f"Scene tensor read body_fields are unknown: {unknown}; " + f"expected={sorted(_TENSOR_BODY_FIELDS)}" + ) + if len(set(fields)) != len(fields): + raise ValueError(f"Scene tensor read body_fields must be unique: {fields}") object.__setattr__(self, "sensor_names", sensors) object.__setattr__(self, "body_names", bodies) + object.__setattr__(self, "body_fields", fields) class SceneTensorReadPlan: @@ -277,6 +335,11 @@ def __init__( entities: Mapping[str, "Entity"], sensor_names: Mapping[str, tuple[str, ...]], body_names: Mapping[str, tuple[str, ...]], + body_fields: Mapping[str, tuple[str, ...]], + packet_names: tuple[str, ...], + packet_sensor_owners: Mapping[str, str], + packet_body_axes: Mapping[str, tuple[str, ...]], + packet_body_blocks: bool, host_plan: HostBridgeTransferPlan | None, ) -> None: self._scene = scene @@ -284,20 +347,24 @@ def __init__( self._entities = MappingProxyType(dict(entities)) self._sensor_names = MappingProxyType(dict(sensor_names)) self._body_names = MappingProxyType(dict(body_names)) + self._body_fields = MappingProxyType(dict(body_fields)) self._host_plan_value = host_plan - self._packet_names = self._aggregate_packet_names() + self._packet_names = packet_names + owners = dict(packet_sensor_owners) + if not packet_body_blocks: + for entity_name, body_names_value in self._body_names.items(): + for body_name in body_names_value: + for field, prefix in _TENSOR_BODY_SENSOR_FIELDS: + if field not in self._body_fields[entity_name]: + continue + name = prefix + body_name + owners.setdefault(name, entity_name) + self._packet_sensor_owners = MappingProxyType(owners) + self._packet_body_axes = MappingProxyType(dict(packet_body_axes)) + self._packet_body_blocks = packet_body_blocks self._aggregate_body_views: dict[str, EntityTensorBodyStateView] = {} self._joint_column_indices: dict[str, tuple[torch.Tensor, torch.Tensor]] = {} - self._packet_sensor_owners: dict[str, str] = {} - self._packet_sensor_owners.update( - (name, entity_name) - for entity_name, names in self._sensor_names.items() - for name in names - ) - for entity_name, body_names_value in self._body_names.items(): - for body_name in body_names_value: - for _, prefix in _TENSOR_BODY_SENSOR_FIELDS: - self._packet_sensor_owners.setdefault(prefix + body_name, entity_name) + self._packet_body_index_cache: dict[tuple[str, tuple[str, ...]], torch.Tensor] = {} self._packet: Mapping[str, torch.Tensor] = MappingProxyType({}) self._refreshed = False self._closed = False @@ -322,6 +389,16 @@ def sensor_names(self) -> Mapping[str, tuple[str, ...]]: def body_names(self) -> Mapping[str, tuple[str, ...]]: return self._body_names + @property + def body_fields(self) -> Mapping[str, tuple[str, ...]]: + """Return the semantic body blocks compiled for each entity.""" + return self._body_fields + + @property + def packet_body_blocks(self) -> bool: + """Return whether the backend publishes native packed body blocks.""" + return self._packet_body_blocks + @property def host_plan(self) -> HostBridgeTransferPlan | None: return self._host_plan_value @@ -336,6 +413,38 @@ def packet_names(self) -> tuple[str, ...]: """Return the immutable backend sensor names compiled into this plan.""" return self._packet_names + def ray_pose_view( + self, entity: "str | Entity", body_names: Sequence[str] | None = None + ) -> SceneTensorPoseView: + """Return a ray-shaped full-batch pose view backed by this phase packet. + The tensors are persistent packed views in the requested body order. + They are valid until the next phase refresh, invalidation, or close. + """ + self._require_open() + owner = self._require_entity(entity) + requested = self._compiled_body_request(owner, body_names) + fields = self._packet_body_fields( + owner, + requested, + required_fields=_TENSOR_POSE_BODY_FIELDS, + capability="ray pose tensor read", + ) + return SceneTensorPoseView( + body_names=requested, + pos_w=fields["pos_w"], + quat_w=fields["quat_w"], + ) + + def packet_body_axis(self, packet_field: str) -> tuple[str, ...]: + """Return immutable aggregate body identity for one packed body field.""" + try: + return self._packet_body_axes[packet_field] + except KeyError as exc: + raise KeyError( + f"Scene tensor body field '{packet_field}' was not compiled into the " + f"packed read; compiled={list(self._packet_body_axes)}" + ) from exc + def refresh(self) -> None: """Publish one full-batch packet for the current read phase.""" self._require_not_closed() @@ -500,14 +609,9 @@ def body_tensor_view( self, entity: "str | Entity", body_names: Sequence[str] | None = None ) -> EntityTensorBodyStateView: """Return entity-local body state from canonical phase sensors.""" + self._require_open() owner = self._require_entity(entity) - if owner._body_names is None: - raise owner._capability_error( - "scene body tensor state", "body_names were not declared in EntityCfg" - ) - requested = self._scene._normalize_tensor_body_request( - owner, tuple(body_names) if body_names is not None else owner._body_names - ) + requested = self._compiled_body_request(owner, body_names) missing = [name for name in requested if name not in set(self._body_names[owner.name])] if missing: raise ValueError( @@ -519,8 +623,31 @@ def body_tensor_view( if aggregate is not None and aggregate.body_names == requested: return aggregate + compiled_fields = self._body_fields[owner.name] + packet_fields_available = all( + _TENSOR_BODY_PACKET_FIELDS[field] in self._packet for field in compiled_fields + ) + if packet_fields_available: + fields = self._packet_body_fields( + owner, + requested, + required_fields=compiled_fields, + capability="aggregate body tensor read", + ) + view = EntityTensorBodyStateView( + body_names=requested, + pos_w=fields.get("pos_w"), + quat_w=fields.get("quat_w"), + lin_vel_w=fields.get("lin_vel_w"), + ang_vel_w=fields.get("ang_vel_w"), + ) + if requested == self._body_names[owner.name]: + self._aggregate_body_views[owner.name] = view + return view capabilities = self._scene._backend.get_tensor_capabilities() - if getattr(capabilities, "tracked_body_views", False): + if getattr(capabilities, "tracked_body_views", False) and set(compiled_fields) == set( + _TENSOR_BODY_FIELDS + ): try: # A single requested body needs no four-block column projection; # the per-body canonical sensors are already the exact packed @@ -553,8 +680,9 @@ def body_tensor_view( self._aggregate_body_views[owner.name] = view return view - fields: dict[str, torch.Tensor] = {} - for field, prefix in _TENSOR_BODY_SENSOR_FIELDS: + stacked_fields: dict[str, torch.Tensor] = {} + for field in compiled_fields: + prefix = _TRACKED_BODY_SENSOR_PREFIXES[field] stacked = [ Entity._validate_tensor_sensor( self._packet[prefix + body_name], @@ -568,15 +696,119 @@ def body_tensor_view( ) for body_name in requested ] - fields[field] = torch.stack(stacked, dim=1) + stacked_fields[field] = torch.stack(stacked, dim=1) return EntityTensorBodyStateView( body_names=requested, - pos_w=fields["pos_w"], - quat_w=fields["quat_w"], - lin_vel_w=fields["lin_vel_w"], - ang_vel_w=fields["ang_vel_w"], + pos_w=stacked_fields["pos_w"], + quat_w=stacked_fields["quat_w"], + lin_vel_w=stacked_fields["lin_vel_w"], + ang_vel_w=stacked_fields["ang_vel_w"], ) + def _compiled_body_request( + self, owner: "Entity", body_names: Sequence[str] | None + ) -> tuple[str, ...]: + if owner._body_names is None: + raise owner._capability_error( + "scene body tensor state", "body_names were not declared in EntityCfg" + ) + return self._scene._normalize_tensor_body_request( + owner, tuple(body_names) if body_names is not None else owner._body_names + ) + + def _packet_body_fields( + self, + owner: "Entity", + requested: tuple[str, ...], + *, + required_fields: tuple[str, ...], + capability: str, + ) -> dict[str, torch.Tensor]: + fields: dict[str, torch.Tensor] = {} + missing: list[str] = [] + compiled_fields = self._body_fields[owner.name] + for field in required_fields: + if field not in compiled_fields: + missing.append(field) + if missing: + raise ValueError( + f"Entity '{owner.name}' {capability} fields {missing} were not compiled " + f"into the packed read; compiled={list(compiled_fields)}" + ) + compiled_bodies = self._body_names[owner.name] + if not self._packet_body_blocks: + canonical_fields: dict[str, torch.Tensor] = {} + for field, prefix in _TENSOR_BODY_SENSOR_FIELDS: + if field not in required_fields: + continue + values = tuple(self._packet.get(prefix + name) for name in requested) + if any(value is None for value in values): + missing.append(field) + continue + canonical_fields[field] = torch.stack( + tuple(value for value in values if value is not None), dim=1 + ) + if missing: + raise ValueError( + f"Entity '{owner.name}' {capability} fields {missing} were not compiled " + f"into the packed read; compiled={list(compiled_fields)}" + ) + return canonical_fields + for field in required_fields: + packet_name = _TENSOR_BODY_PACKET_FIELDS[field] + raw = self._packet.get(packet_name) + if raw is None: + missing.append(field) + continue + axis = self._packet_body_axes.get(packet_name, ()) + if not axis: + missing.append(field) + continue + width = _TENSOR_PACKET_FIELD_WIDTHS[packet_name] + expected = (self._scene._backend.num_envs, len(axis) * width) + if tuple(raw.shape) != expected: + raise ValueError( + f"Entity '{owner.name}' {capability} field '{field}' has shape " + f"{tuple(raw.shape)}; expected {expected}" + ) + if raw.dtype != torch.float32 or raw.device != self._device: + raise TypeError( + f"Entity '{owner.name}' {capability} field '{field}' must be float32 " + f"on {self._device}; got {raw.dtype} on {raw.device}" + ) + # ``unflatten`` preserves a view for row-strided packet slices. + # ``reshape`` would silently materialize a contiguous copy when the + # backend publishes one packed body block beside unrelated columns. + value = raw.unflatten(1, (len(axis), width)) + owner_axis = tuple(_packet_body_name(owner.name, name) for name in compiled_bodies) + requested_axis = tuple(_packet_body_name(owner.name, name) for name in requested) + if requested_axis != owner_axis or owner_axis != axis: + rows = self._packet_body_indices(requested_axis, axis) + value = value.index_select(1, rows) + fields[field] = value + if missing: + raise ValueError( + f"Entity '{owner.name}' {capability} fields {missing} were not compiled " + f"into the packed read; compiled={list(self._body_fields.get(owner.name, ()))}" + ) + return fields + + def _packet_body_indices( + self, requested: tuple[str, ...], compiled: tuple[str, ...] + ) -> torch.Tensor: + key: tuple[str, tuple[str, ...]] = ("body", compiled) + cached = self._packet_body_index_cache.get(key) + if cached is not None: + return cached + positions = {name: index for index, name in enumerate(compiled)} + value = torch.tensor( + [positions[name] for name in requested], + dtype=torch.int64, + device=self._device, + ) + self._packet_body_index_cache[key] = value + return value + def _validated_aggregate_body_view( self, owner: "Entity", views: Any ) -> EntityTensorBodyStateView: @@ -632,18 +864,32 @@ def close(self) -> None: self._host_plan_value.close() finally: self._aggregate_body_views.clear() + self._packet_body_index_cache.clear() self._packet = MappingProxyType({}) self._refreshed = False self._closed = True - def _aggregate_packet_names(self) -> tuple[str, ...]: - names: list[str] = [] - for entity_names in self._sensor_names.values(): - names.extend(entity_names) - for entity_names in self._body_names.values(): - for body_name in entity_names: - names.extend(prefix + body_name for _, prefix in _TENSOR_BODY_SENSOR_FIELDS) - return tuple(dict.fromkeys(names)) + @staticmethod + def _aggregate_packet_body_axes( + body_names: Mapping[str, tuple[str, ...]], + body_fields: Mapping[str, set[str] | tuple[str, ...]], + ) -> dict[str, tuple[str, ...]]: + axes: dict[str, tuple[str, ...]] = {} + for entity_name, names in body_names.items(): + if not names: + continue + for field in body_fields[entity_name]: + packet_name = _TENSOR_BODY_PACKET_FIELDS[field] + axis = axes.setdefault(packet_name, ()) + axes[packet_name] = ( + *axis, + *( + _packet_body_name(entity_name, body_name) + for body_name in names + if _packet_body_name(entity_name, body_name) not in axis + ), + ) + return axes def _read_packet( self, reader: Callable[[], Mapping[str, Any]] | None @@ -663,9 +909,23 @@ def _read_packet( # stable adapter-owned storage here; only a fresh call proves # that storage reflects the current post-step/post-reset state. raw = dict(backend.get_state_views(("qpos", "qvel"), device=self._device)) + body_packet_names = tuple( + name for name in self._packet_names if name not in self._packet_sensor_owners + ) + body_view: Any = None + if body_packet_names: + body_names = tuple( + name.partition("/")[2] if "/" in name else name + for name in self._packet_body_axes[body_packet_names[0]] + ) + body_view = backend.get_tracked_body_views(body_names, device=self._device) for name in self._packet_names: - view = backend.get_sensor_view(name, device=self._device) - raw[name] = view + if name in self._packet_sensor_owners: + raw[name] = backend.get_sensor_view(name, device=self._device) + else: + raw[name] = getattr(body_view, _PACKET_FIELD_BODY_FIELDS[name]).reshape( + backend.num_envs, -1 + ) except (AttributeError, KeyError, TypeError, ValueError, NotImplementedError) as exc: raise type(exc)( "Manager scene device-resident tensor read on backend " @@ -706,17 +966,38 @@ def _publish_packet(self, packet: dict[str, torch.Tensor]) -> None: f"got {value.dtype} on {value.device}" ) for name in self._packet_names: - owner = self._packet_sensor_owners[name] - packet[name] = Entity._validate_tensor_sensor( - packet[name], - entity_name=owner, - backend_type=backend.backend_type, - sensor_name=name, - expected_width=None, - num_envs=num_envs, - device=self._device, - finite=False, + value = packet[name] + if name in self._packet_sensor_owners: + packet[name] = Entity._validate_tensor_sensor( + value, + entity_name=self._packet_sensor_owners[name], + backend_type=backend.backend_type, + sensor_name=name, + expected_width=None, + num_envs=num_envs, + device=self._device, + finite=False, + ) + continue + if not isinstance(value, torch.Tensor): + raise TypeError( + f"Scene tensor body field '{name}' on backend " + f"'{backend.backend_type}' is {type(value).__name__}, expected torch.Tensor" + ) + expected = ( + num_envs, + len(self._packet_body_axes[name]) * _TENSOR_PACKET_FIELD_WIDTHS[name], ) + if tuple(value.shape) != expected: + raise ValueError( + f"Scene tensor body field '{name}' on backend " + f"'{backend.backend_type}' has shape {tuple(value.shape)}; expected {expected}" + ) + if value.dtype != torch.float32 or value.device != self._device: + raise TypeError( + f"Scene tensor body field '{name}' must be float32 on {self._device}; " + f"got {value.dtype} on {value.device}" + ) self._packet = MappingProxyType(packet) self._aggregate_body_views.clear() self._refreshed = True @@ -4224,6 +4505,7 @@ def compile_tensor_reads( entities: dict[str, Entity] = {} per_entity_sensors: dict[str, tuple[str, ...]] = {} per_entity_bodies: dict[str, tuple[str, ...]] = {} + per_entity_body_fields: dict[str, set[str]] = {} request_owners: dict[str, str] = {} for spec in normalized: try: @@ -4236,6 +4518,7 @@ def compile_tensor_reads( entities[spec.entity] = entity sensors = per_entity_sensors.setdefault(spec.entity, ()) bodies = per_entity_bodies.setdefault(spec.entity, ()) + fields = per_entity_body_fields.setdefault(spec.entity, set()) for name in spec.sensor_names: owner = request_owners.get(name) if owner is not None and owner != spec.entity: @@ -4246,6 +4529,8 @@ def compile_tensor_reads( request_owners[name] = spec.entity per_entity_sensors[spec.entity] = (*sensors, *spec.sensor_names) per_entity_bodies[spec.entity] = (*bodies, *spec.body_names) + if spec.body_names: + fields.update(spec.body_fields) expanded_bodies: dict[str, tuple[str, ...]] = {} for entity_name, body_names in per_entity_bodies.items(): @@ -4265,17 +4550,69 @@ def compile_tensor_reads( capabilities = self._backend.get_tensor_capabilities() self._validate_tensor_read_capabilities(capabilities, resolved_device) - sensor_names = self._aggregate_tensor_read_names( - entities, per_entity_sensors, expanded_bodies - ) - host_plan: HostBridgeTransferPlan | None = None - if capabilities.execution is TensorExecution.HOST_BRIDGE: - spec = TensorIOSpec( - state_fields=("qpos", "qvel"), - sensor_names=sensor_names, - device=resolved_device, + host_body_blocks = False + host_sensor_packet_names: tuple[str, ...] | None = None + sensor_packet_names = self._aggregate_tensor_read_names(tuple(entities), per_entity_sensors) + requested_body_fields = any(expanded_bodies.values()) + packet_body_blocks = False + if capabilities.execution is TensorExecution.DEVICE_RESIDENT and requested_body_fields: + packet_body_blocks = bool(getattr(capabilities, "tracked_body_views", False)) + packet_names = self._aggregate_tensor_read_names( + tuple(entities), per_entity_sensors, expanded_bodies, per_entity_body_fields + ) + packet_body_axes: dict[str, tuple[str, ...]] = {} + packet_body_names = tuple( + _TENSOR_BODY_PACKET_FIELDS[field] + for field in _TENSOR_BODY_FIELD_ORDER + if any(field in fields for fields in per_entity_body_fields.values()) + ) + if packet_body_blocks: + packet_body_axes = SceneTensorReadPlan._aggregate_packet_body_axes( + expanded_bodies, per_entity_body_fields ) + packet_names = (*sensor_packet_names, *packet_body_names) + if capabilities.execution is TensorExecution.HOST_BRIDGE: + host_body_blocks = {"body_names", "body_fields"}.issubset( + inspect.signature(TensorIOSpec).parameters + ) + host_pose_only_blocks = host_body_blocks and all( + not (set(fields) - set(_TENSOR_POSE_BODY_FIELDS)) + for fields in per_entity_body_fields.values() + ) + if host_pose_only_blocks and any(expanded_bodies.values()): + packet_body_blocks = True + packet_names = (*sensor_packet_names, *packet_body_names) + host_sensor_packet_names = sensor_packet_names + for entity_name, body_names_value in expanded_bodies.items(): + if not body_names_value: + continue + for field in per_entity_body_fields[entity_name]: + packet_name = _TENSOR_BODY_PACKET_FIELDS[field] + axis = packet_body_axes.setdefault(packet_name, ()) + packet_body_axes[packet_name] = ( + *axis, + *( + _packet_body_name(entity_name, body_name) + for body_name in body_names_value + if _packet_body_name(entity_name, body_name) not in axis + ), + ) + tensor_io_options: dict[str, Any] = { + "state_fields": ("qpos", "qvel"), + "sensor_names": ( + sensor_packet_names if host_sensor_packet_names is not None else packet_names + ), + "device": resolved_device, + } + if host_pose_only_blocks: + tensor_io_options.update( + { + "body_names": self._aggregate_tensor_read_body_names(expanded_bodies), + "body_fields": ("pos_w", "quat_w"), + } + ) + spec = TensorIOSpec(**tensor_io_options) try: host_plan = self._backend.bind_host_bridge_io(spec) except (AttributeError, KeyError, TypeError, ValueError, NotImplementedError) as exc: @@ -4290,6 +4627,18 @@ def compile_tensor_reads( entities=entities, sensor_names=per_entity_sensors, body_names=expanded_bodies, + body_fields={ + name: tuple(field for field in _TENSOR_BODY_FIELD_ORDER if field in fields) + for name, fields in per_entity_body_fields.items() + }, + packet_names=packet_names, + packet_sensor_owners={ + name: entity_name + for entity_name, names in per_entity_sensors.items() + for name in names + }, + packet_body_axes=packet_body_axes, + packet_body_blocks=packet_body_blocks, host_plan=host_plan, ) @@ -4343,9 +4692,10 @@ def _validate_tensor_read_capabilities( @staticmethod def _aggregate_tensor_read_names( - entities: Mapping[str, "Entity"], + entity_names: Sequence[str], sensor_names: Mapping[str, tuple[str, ...]], - body_names: Mapping[str, tuple[str, ...]], + body_names: Mapping[str, tuple[str, ...]] | None = None, + body_fields: Mapping[str, set[str] | tuple[str, ...]] | None = None, ) -> tuple[str, ...]: aggregate: list[str] = [] @@ -4355,13 +4705,32 @@ def add(name: str) -> None: aggregate.append(name) seen: set[str] = set() - for entity_name in entities: + for entity_name in entity_names: for name in sensor_names[entity_name]: add(name) - for entity_name in entities: - for body_name in body_names[entity_name]: - for _, prefix in _TENSOR_BODY_SENSOR_FIELDS: - add(prefix + body_name) + if body_names is not None and body_fields is not None: + for entity_name in entity_names: + fields = body_fields[entity_name] + if entity_name not in body_names: + continue + for body_name in body_names[entity_name]: + for field, prefix in _TENSOR_BODY_SENSOR_FIELDS: + if field in fields: + add(prefix + body_name) + return tuple(aggregate) + + @staticmethod + def _aggregate_tensor_read_body_names( + body_names: Mapping[str, tuple[str, ...]], + ) -> tuple[str, ...]: + aggregate: list[str] = [] + for entity_name, names in body_names.items(): + if not names: + continue + for name in names: + packet_name = _packet_body_name(entity_name, name) + if packet_name not in aggregate: + aggregate.append(packet_name) return tuple(aggregate) @contextmanager @@ -4466,5 +4835,6 @@ def __len__(self) -> int: "EntityTensorSensorViews", "EntityTensorStateView", "SceneTensorReadPlan", + "SceneTensorPoseView", "SceneTensorReadSpec", ] diff --git a/src/unilab/envs/manager_based_rl_env.py b/src/unilab/envs/manager_based_rl_env.py index 9bd1ae60b..66aa466f6 100644 --- a/src/unilab/envs/manager_based_rl_env.py +++ b/src/unilab/envs/manager_based_rl_env.py @@ -427,12 +427,14 @@ def __init__(self, cfg: ManagerBasedRlEnvCfg, backend: SimBackend, num_envs: int env_ids=self._all_env_rows if self.event_manager.uses_startup_rows else None, ) self._compile_tensor_read_plan() + self.observation_manager.bind_read_phase() self.command_manager.bind_read_phase() self._validate_manager_tensor_runtime() def _compile_tensor_read_plan(self) -> None: """Compile the scene's only packed tensor read phase.""" specs: list[SceneTensorReadSpec] = [] + specs.extend(self._observation_ray_read_specs()) specs.extend(self._action_tensor_read_specs()) specs.extend(self._observation_tensor_read_specs()) specs.extend(self._task_tensor_state_read_specs()) @@ -507,6 +509,10 @@ def _observation_tensor_read_specs(self) -> list[SceneTensorReadSpec]: ) ) term = term_cfg.func + if getattr(term, "ray_pattern", None) is not None: + # Ray terms already declared a pose-only body request above. + # They must not also add the legacy four-field body-state set. + continue body_names = term_cfg.params.get("tensor_body_names") if body_names is None: body_names = getattr(term, "tensor_body_names", None) @@ -532,6 +538,30 @@ def _observation_tensor_read_specs(self) -> list[SceneTensorReadSpec]: specs.append(SceneTensorReadSpec(entity=entity_name, body_names=tuple(body_names))) return specs + def _observation_ray_read_specs(self) -> list[SceneTensorReadSpec]: + """Collect pose-only blocks needed by tensor ray observations.""" + specs: list[SceneTensorReadSpec] = [] + for group_name, terms in self.observation_manager.active_terms.items(): + for name in terms: + term_cfg = self.observation_manager.get_term_cfg(group_name, name) + ray_pattern = getattr(term_cfg.func, "ray_pattern", None) + if ray_pattern is None: + continue + if not isinstance(ray_pattern, mdp.RayPatternCfg): + raise TypeError( + "ManagerBasedRlEnv ray declaration for observation term " + f"'{group_name}/{name}' must be RayPatternCfg; got " + f"{type(ray_pattern).__name__}" + ) + specs.append( + SceneTensorReadSpec( + entity=ray_pattern.entity_name, + body_names=tuple(ray_pattern.body_names), + body_fields=("pos_w", "quat_w"), + ) + ) + return specs + def _manager_term_tensor_read_specs(self) -> list[SceneTensorReadSpec]: specs: list[SceneTensorReadSpec] = [] narrow = ( @@ -1725,6 +1755,7 @@ def import_training_state(self, state: Mapping[str, Any]) -> None: self._sim_step_counter = self.step_counter * self._cfg.sim_substeps def close(self) -> None: + self.observation_manager.close() read_plan = self.scene.read_plan if read_plan is not None: read_plan.close() diff --git a/src/unilab/envs/mdp/__init__.py b/src/unilab/envs/mdp/__init__.py index a7e460a0b..92fcabe16 100644 --- a/src/unilab/envs/mdp/__init__.py +++ b/src/unilab/envs/mdp/__init__.py @@ -58,6 +58,8 @@ from unilab.envs.mdp.observations import ( projected_gravity_imu_misaligned as projected_gravity_imu_misaligned, ) +from unilab.envs.mdp.ray_observation import RayPatternCfg as RayPatternCfg +from unilab.envs.mdp.ray_observation import RayScan as RayScan from unilab.envs.mdp.recorders import LifecycleCounterRecorder as LifecycleCounterRecorder from unilab.envs.mdp.rewards import action_acc_l2 as action_acc_l2 from unilab.envs.mdp.rewards import action_rate_l2 as action_rate_l2 @@ -96,6 +98,8 @@ "UniformVelocityCommandCfg", "UniformPoseCommand", "UniformPoseCommandCfg", + "RayPatternCfg", + "RayScan", "action_acc_l2", "action_rate_l2", "apply_body_impulse", diff --git a/src/unilab/envs/mdp/ray_observation.py b/src/unilab/envs/mdp/ray_observation.py new file mode 100644 index 000000000..27d39c5c9 --- /dev/null +++ b/src/unilab/envs/mdp/ray_observation.py @@ -0,0 +1,197 @@ +"""Tensor-native ray-pattern observation terms for the Manager runtime.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Sequence + +import torch +from unisim.backend.base import torch_device_matches + +from unilab.managers.manager_base import ManagerTermBase, ManagerTermBaseCfg + +if TYPE_CHECKING: + from unilab.managers.types import ManagerBasedRlEnv + + +@dataclass(frozen=True) +class RayPatternCfg: + """Declare one observation-owned ray pattern. + + ``caster`` is a public UniSim ray-caster factory name, not an imported + plugin class. ``body_names`` is entity-local and selects both the backend + ray-scene projection and the pose-only packed read block. + """ + + entity_name: str + body_names: Sequence[str] + num_rays: int + pattern_start: float + pattern_stop: float + max_distance: float + caster: str = "uni_ray" + pattern_center_height: float = 0.0 + caster_options: dict[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not isinstance(self.entity_name, str) or not self.entity_name: + raise TypeError("RayPatternCfg entity_name must be a non-empty string") + if isinstance(self.body_names, (str, bytes)): + raise TypeError("RayPatternCfg body_names must be a sequence of strings") + names = tuple(self.body_names) + if not names or any(not isinstance(name, str) or not name for name in names): + raise TypeError("RayPatternCfg body_names must be non-empty strings") + if len(set(names)) != len(names): + raise ValueError(f"RayPatternCfg body_names must be unique: {names}") + if isinstance(self.num_rays, bool) or not isinstance(self.num_rays, int): + raise TypeError("RayPatternCfg num_rays must be an integer") + if self.num_rays <= 0: + raise ValueError("RayPatternCfg num_rays must be positive") + for name, value in ( + ("pattern_start", self.pattern_start), + ("pattern_stop", self.pattern_stop), + ("max_distance", self.max_distance), + ("pattern_center_height", self.pattern_center_height), + ): + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise TypeError(f"RayPatternCfg {name} must be a real number") + if not math.isfinite(float(value)): + raise ValueError(f"RayPatternCfg {name} must be finite") + if self.max_distance <= 0.0: + raise ValueError("RayPatternCfg max_distance must be positive") + if not isinstance(self.caster, str) or not self.caster: + raise TypeError("RayPatternCfg caster must be a non-empty factory name") + if not isinstance(self.caster_options, dict): + raise TypeError("RayPatternCfg caster_options must be a mapping") + object.__setattr__(self, "body_names", names) + + +class RayScan(ManagerTermBase): + """Bind one tensor caster and publish its distance observation carrier.""" + + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): + super().__init__(env) + pattern = cfg.params.get("ray_pattern") + if not isinstance(pattern, RayPatternCfg): + raise TypeError( + "RayScan requires params['ray_pattern'] to be a RayPatternCfg; " + f"got {type(pattern).__name__}" + ) + self.ray_pattern = pattern + self.tensor_body_names = tuple(pattern.body_names) + self.observation_width = pattern.num_rays + self._device = torch.device(env.device) + self._entity_name = pattern.entity_name + self._body_names = tuple(pattern.body_names) + self._caster: Any = None + self._bound = False + self._closed = False + self._distance: torch.Tensor | None = None + self._origins: torch.Tensor | None = None + self._directions: torch.Tensor | None = None + + def bind_read_phase(self) -> None: + """Materialize and bind the caster after packed pose reads compile.""" + if self._bound or self._closed: + return + read_plan = getattr(self._env.scene, "read_plan", None) + if read_plan is None: + raise NotImplementedError("RayScan requires the Manager packed Torch read plan") + try: + scene = self._env.backend.get_ray_scene_description(self._body_names) + caster = _create_ray_caster( + self.ray_pattern.caster, + num_envs=self._env.num_envs, + num_rays=self.ray_pattern.num_rays, + device=self._device, + **self.ray_pattern.caster_options, + ) + capabilities = caster.get_ray_capabilities() + if not ( + capabilities.supports_tensor_pose_sync + and capabilities.supports_tensor_rays + and capabilities.supports_device_output + and capabilities.stream_event_ownership == "caller" + and torch_device_matches(capabilities.torch_devices, self._device) + ): + raise NotImplementedError( + f"RayScan caster '{self.ray_pattern.caster}' does not support the " + f"tensor pose/ray/output profile on {self._device}" + ) + caster.materialize(scene) + caster.bind_pose_layout(tuple(scene.body_names)) + except (AttributeError, KeyError, TypeError, ValueError, NotImplementedError) as exc: + raise type(exc)( + f"RayScan caster '{self.ray_pattern.caster}' read-phase binding failed: {exc}" + ) from exc + self._caster = caster + self._distance = torch.empty( + (self._env.num_envs, self.ray_pattern.num_rays), + dtype=torch.float32, + device=self._device, + ) + angles = torch.linspace( + self.ray_pattern.pattern_start, + self.ray_pattern.pattern_stop, + self.ray_pattern.num_rays, + dtype=torch.float32, + device=self._device, + ) + self._directions = torch.stack( + (torch.cos(angles), torch.sin(angles), torch.zeros_like(angles)), dim=-1 + ).contiguous() + self._origins = torch.zeros( + (self._env.num_envs, self.ray_pattern.num_rays, 3), + dtype=torch.float32, + device=self._device, + ) + self._origins[:, :, 2] = self.ray_pattern.pattern_center_height + self._origins = self._origins.contiguous() + self._bound = True + + def __call__(self, env: ManagerBasedRlEnv) -> torch.Tensor: + if not self._bound or self._caster is None or self._distance is None: + raise NotImplementedError("RayScan must bind its read phase before compute") + read_plan = getattr(env.scene, "read_plan", None) + if read_plan is None or not read_plan.ready: + raise RuntimeError("RayScan requires a refreshed packed read phase") + pose = read_plan.ray_pose_view(self._entity_name, self._body_names) + try: + self._caster.update_pose_tensor(pose) + result = self._caster.trace_tensor( + self._origins, + self._directions, + self.ray_pattern.max_distance, + ) + except (AttributeError, KeyError, TypeError, ValueError, NotImplementedError) as exc: + raise type(exc)(f"RayScan tensor trace failed: {exc}") from exc + self._distance.copy_(result.distance) + return self._distance + + def close(self) -> None: + """Release caster resources before the scene read plan closes.""" + self._closed = True + caster, self._caster = self._caster, None + if caster is not None: + caster.close() + + def reset(self, env_ids: torch.Tensor | slice | None = None) -> None: + """Retain caster state; observations refresh from the next read phase.""" + del env_ids + + +__all__ = ["RayPatternCfg", "RayScan"] + + +def _create_ray_caster(caster: str, **options: Any) -> Any: + """Lazily call the installed UniSim public caster factory.""" + import unisim + + factory = getattr(unisim, "create_ray_caster", None) + if not callable(factory): + raise NotImplementedError( + "installed unisim-core does not expose create_ray_caster; use the " + "ray-local optional dependency profile" + ) + return factory(caster, **options) diff --git a/src/unilab/managers/observation_manager.py b/src/unilab/managers/observation_manager.py index 37222b72f..fbae777fa 100644 --- a/src/unilab/managers/observation_manager.py +++ b/src/unilab/managers/observation_manager.py @@ -362,6 +362,23 @@ def reset(self, env_ids: torch.Tensor | slice | None = None) -> dict[str, float] mod.reset(reset_mask) return {} + def bind_read_phase(self) -> None: + """Bind class terms after the scene's packed tensor read plan exists.""" + for group_cfg in self._group_obs_class_term_cfgs.values(): + for term_cfg in group_cfg: + binder = getattr(term_cfg.func, "bind_read_phase", None) + if not callable(binder): + continue + binder() + + def close(self) -> None: + """Forward close to class terms that own external resources.""" + for group_cfg in self._group_obs_class_term_cfgs.values(): + for term_cfg in group_cfg: + closer = getattr(term_cfg.func, "close", None) + if callable(closer): + closer() + def _check_and_handle_nans( self, tensor: torch.Tensor, @@ -787,6 +804,16 @@ def _observations_to_tensor_boundary( return values raise TypeError("ObservationManager results must be torch.Tensor carriers") + def _deferred_term_width(self, term_cfg: ObservationTermCfg) -> int: + """Resolve the fixed width of a term that binds after the read phase.""" + width = getattr(term_cfg.func, "observation_width", None) + if isinstance(width, bool) or not isinstance(width, int) or width <= 0: + raise TypeError( + f"ObservationManager term '{term_cfg.func}' with deferred read binding " + "must declare a positive integer observation_width" + ) + return width + def _prepare_terms(self) -> None: self._group_obs_term_names: dict[str, list[str]] = dict() self._group_obs_term_dim: dict[str, list[tuple[int, ...]]] = dict() @@ -878,30 +905,34 @@ def _prepare_terms(self) -> None: if hasattr(term_cfg.func, "reset") and callable(term_cfg.func.reset): self._group_obs_class_term_cfgs[group_name].append(term_cfg) - initial_obs = term_cfg.func(self._env, **term_cfg.params) - if not isinstance(initial_obs, torch.Tensor): - raise TypeError( - f"ObservationManager term '{group_name}/{term_name}' returned " - f"{type(initial_obs).__name__}, expected torch.Tensor." - ) - if initial_obs.dtype != torch.float32 or initial_obs.device != self._device: - raise TypeError( - f"ObservationManager term '{group_name}/{term_name}' must return " - f"float32 observations on {self._device}; got " - f"{initial_obs.dtype} on {initial_obs.device}." - ) - if initial_obs.ndim < 2 or initial_obs.shape[0] != self.num_envs: - raise ValueError( - f"ObservationManager term '{group_name}/{term_name}' returned shape " - f"{initial_obs.shape}, expected (num_envs, ...) with " - f"num_envs={self.num_envs}." - ) - obs_dims = tuple(initial_obs.shape) + deferred_bind = callable(getattr(term_cfg.func, "bind_read_phase", None)) + if not deferred_bind: + initial_obs = term_cfg.func(self._env, **term_cfg.params) + if not isinstance(initial_obs, torch.Tensor): + raise TypeError( + f"ObservationManager term '{group_name}/{term_name}' returned " + f"{type(initial_obs).__name__}, expected torch.Tensor." + ) + if initial_obs.dtype != torch.float32 or initial_obs.device != self._device: + raise TypeError( + f"ObservationManager term '{group_name}/{term_name}' must return " + f"float32 observations on {self._device}; got " + f"{initial_obs.dtype} on {initial_obs.device}." + ) + if initial_obs.ndim < 2 or initial_obs.shape[0] != self.num_envs: + raise ValueError( + f"ObservationManager term '{group_name}/{term_name}' returned shape " + f"{initial_obs.shape}, expected (num_envs, ...) with " + f"num_envs={self.num_envs}." + ) + obs_dims = tuple(initial_obs.shape) + else: + obs_dims = (self.num_envs, self._deferred_term_width(term_cfg)) if term_cfg.scale is not None: scale = torch.as_tensor(term_cfg.scale, dtype=torch.float32) self._scale_tensors[id(term_cfg)] = ( - scale.broadcast_to(initial_obs.shape).clone().to(device=self._device) + scale.broadcast_to(obs_dims).clone().to(device=self._device) ) if term_cfg.noise is not None and isinstance( diff --git a/src/unilab/managers/types.py b/src/unilab/managers/types.py index 55c0a8610..421bb5165 100644 --- a/src/unilab/managers/types.py +++ b/src/unilab/managers/types.py @@ -292,6 +292,9 @@ def action_manager(self) -> ManagerActionManager: ... @property def command_manager(self) -> ManagerCommandManager: ... + @property + def observation_manager(self) -> Any: ... + @property def event_manager(self) -> ManagerEventManager: ... diff --git a/tests/base/test_entity_tensor_read_plan.py b/tests/base/test_entity_tensor_read_plan.py index 550b87742..01f1b2b6f 100644 --- a/tests/base/test_entity_tensor_read_plan.py +++ b/tests/base/test_entity_tensor_read_plan.py @@ -2,6 +2,7 @@ from __future__ import annotations +from dataclasses import dataclass, field from types import MethodType from typing import cast @@ -36,6 +37,7 @@ def __init__(self, *, device_resident: bool = False) -> None: self.sensor_view_calls = 0 self.aggregate_body_view_calls = 0 self.closes = 0 + self.io_specs: list[TensorIOSpec] = [] self.qpos = torch.tensor([[0.1], [0.3]], dtype=torch.float32) self.qvel = torch.tensor([[10.0], [20.0]], dtype=torch.float32) self.sensors = { @@ -49,6 +51,12 @@ def __init__(self, *, device_resident: bool = False) -> None: "track_linvel_w_hip": torch.zeros((2, 3), dtype=torch.float32), "track_angvel_w_hip": torch.ones((2, 3), dtype=torch.float32), } + self.body_blocks = { + "body_pos_w": self.sensors["track_pos_w_hip"].reshape(2, 3).contiguous(), + "body_quat_w": self.sensors["track_quat_w_hip"].reshape(2, 4).contiguous(), + "body_lin_vel_w": self.sensors["track_linvel_w_hip"].reshape(2, 3).contiguous(), + "body_ang_vel_w": self.sensors["track_angvel_w_hip"].reshape(2, 3).contiguous(), + } self._selected_rows = torch.empty((0,), dtype=torch.int64) def get_tensor_capabilities(self) -> TensorLifecycleCapabilities: @@ -97,6 +105,7 @@ def get_sensor_view(self, name: str, device: str | torch.device = "cpu") -> torc def get_tracked_body_views(self, body_names=None, device=None) -> TrackedBodyStateViews: self.aggregate_body_view_calls += 1 names = ("hip",) if body_names is None else tuple(body_names) + names = tuple(name.rsplit("/", 1)[-1] for name in names) if names != ("hip",): raise ValueError(f"unknown tracked bodies: {names}") target = torch.device(device) if device is not None else torch.device("cpu") @@ -183,6 +192,8 @@ def get_joint_dof_vel_indices(self, names) -> np.ndarray: def bind_host_bridge_io(self, spec: TensorIOSpec): backend = self + if getattr(spec, "body_names", ()): + backend.io_specs.append(spec) class _Plan(HostBridgeTransferPlan): def __init__(self, plan_spec: TensorIOSpec) -> None: @@ -226,6 +237,9 @@ def close(self) -> None: def _packet(self) -> dict[str, torch.Tensor]: packet = self.get_state_views(("qpos", "qvel"), device="cpu") packet.update(self.sensors) + if self.device_resident: + return packet + packet.update(self.body_blocks) return packet @@ -242,6 +256,18 @@ def _scene(backend: _SceneTensorBackend) -> EntityScene: ) +def _use_body_capable_tensor_spec(monkeypatch: pytest.MonkeyPatch) -> None: + """Expose a UniSim-2.2-shaped public TensorIOSpec for packed-block tests.""" + + @dataclass(frozen=True) + class _BodyCapableTensorIOSpec(TensorIOSpec): + body_names: tuple[str, ...] = () + body_fields: tuple[str, ...] = () + device: object | None = field(default=None) + + monkeypatch.setattr("unilab.base.entity.TensorIOSpec", _BodyCapableTensorIOSpec) + + def _specs() -> tuple[SceneTensorReadSpec, ...]: return ( SceneTensorReadSpec(entity="robot", sensor_names=("imu_gyro",)), @@ -249,8 +275,11 @@ def _specs() -> tuple[SceneTensorReadSpec, ...]: ) -def test_host_bridge_scene_plan_packs_all_reads_into_one_phase_transfer() -> None: +def test_host_bridge_scene_plan_packs_all_reads_into_one_phase_transfer( + monkeypatch: pytest.MonkeyPatch, +) -> None: backend = _SceneTensorBackend() + _use_body_capable_tensor_spec(monkeypatch) scene = _scene(backend) plan = scene.compile_tensor_reads("cpu", _specs()) @@ -276,6 +305,29 @@ def test_host_bridge_scene_plan_packs_all_reads_into_one_phase_transfer() -> Non } plan.close() + + +def test_scene_plan_pose_reads_fall_back_to_canonical_sensors_on_older_host_contracts() -> None: + backend = _SceneTensorBackend() + scene = _scene(backend) + plan = scene.compile_tensor_reads( + "cpu", + ( + SceneTensorReadSpec( + entity="robot", body_names=("hip",), body_fields=("pos_w", "quat_w") + ), + ), + ) + + assert plan.packet_body_blocks is False + assert plan.packet_names == ("track_pos_w_hip", "track_quat_w_hip") + assert backend.io_specs == [] + plan.refresh() + pose = plan.ray_pose_view("robot") + assert pose.body_names == ("hip",) + torch.testing.assert_close(pose.pos_w[:, 0], backend.sensors["track_pos_w_hip"]) + torch.testing.assert_close(pose.quat_w[:, 0], backend.sensors["track_quat_w_hip"]) + plan.close() assert backend.closes == 1 @@ -331,7 +383,7 @@ def test_device_resident_scene_plan_refreshes_every_sensor_boundary() -> None: torch.testing.assert_close(second, torch.full((2, 3), 7.0)) # Every compiled sensor crosses its public backend boundary on every phase. # Genesis/Newton derived sensors are stable outputs, not native-live arrays. - assert backend.sensor_view_calls == calls_after_first_refresh + len(plan.packet_names) + assert backend.sensor_view_calls == calls_after_first_refresh + len(plan.sensor_names["robot"]) assert plan.transfer_stats == {} assert plan.last_timing == {} plan.close() @@ -350,7 +402,9 @@ def test_device_resident_selected_reset_refreshes_body_sensor_boundary() -> None body = plan.body_tensor_view("robot") torch.testing.assert_close(body.pos_w[:, 0], backend.sensors["track_pos_w_hip"]) - assert backend.sensor_view_calls == calls_after_selected_refresh + len(plan.packet_names) + assert backend.sensor_view_calls == calls_after_selected_refresh + len( + plan.sensor_names["robot"] + ) plan.close() @@ -385,6 +439,105 @@ def counting_qpos_indices(names): plan.close() +def test_scene_plan_pose_reads_use_packed_pose_blocks_without_velocities( + monkeypatch: pytest.MonkeyPatch, +) -> None: + backend = _SceneTensorBackend() + _use_body_capable_tensor_spec(monkeypatch) + scene = _scene(backend) + plan = scene.compile_tensor_reads( + "cpu", + ( + SceneTensorReadSpec( + entity="robot", body_names=("hip",), body_fields=("pos_w", "quat_w") + ), + ), + ) + + assert plan.body_fields == {"robot": ("pos_w", "quat_w")} + assert plan.packet_names == ("body_pos_w", "body_quat_w") + if backend.io_specs: + assert backend.io_specs[0].body_names == ("robot/hip",) + assert backend.io_specs[0].body_fields == ("pos_w", "quat_w") + assert plan.packet_body_axis("body_pos_w") == ("robot/hip",) + plan.refresh() + + pose = plan.ray_pose_view("robot") + assert pose.body_names == ("hip",) + assert pose.pos_w.shape == (2, 1, 3) + assert pose.quat_w.shape == (2, 1, 4) + torch.testing.assert_close(pose.pos_w[:, 0], backend.sensors["track_pos_w_hip"]) + torch.testing.assert_close(pose.quat_w[:, 0], backend.sensors["track_quat_w_hip"]) + body = plan.body_tensor_view("robot") + assert body.pos_w is not None and body.quat_w is not None + assert body.lin_vel_w is None and body.ang_vel_w is None + plan.close() + + +def test_scene_plan_ray_pose_view_preserves_row_strided_packet_slices( + monkeypatch: pytest.MonkeyPatch, +) -> None: + backend = _SceneTensorBackend() + _use_body_capable_tensor_spec(monkeypatch) + scene = _scene(backend) + plan = scene.compile_tensor_reads( + "cpu", + ( + SceneTensorReadSpec( + entity="robot", body_names=("hip",), body_fields=("pos_w", "quat_w") + ), + ), + ) + + storage = torch.zeros((2, 12), dtype=torch.float32) + pos = torch.arange(6, dtype=torch.float32).reshape(2, 3) + quat = torch.tensor(((1.0, 0.0, 0.0, 0.0), (1.0, 0.0, 0.0, 0.0)), dtype=torch.float32) + storage[:, 3:6] = pos + storage[:, 7:11] = quat + backend.body_blocks["body_pos_w"] = storage[:, 3:6] + backend.body_blocks["body_quat_w"] = storage[:, 7:11] + plan.refresh() + pose = plan.ray_pose_view("robot") + + assert not pose.pos_w.is_contiguous() + assert not pose.quat_w.is_contiguous() + assert pose.pos_w.data_ptr() == storage[:, 3:6].data_ptr() + assert pose.quat_w.data_ptr() == storage[:, 7:11].data_ptr() + assert pose.pos_w.stride() == (12, 3, 1) + assert pose.quat_w.stride() == (12, 4, 1) + storage[0, 3] = 42.0 + storage[1, 7] = -1.0 + assert pose.pos_w[0, 0, 0].item() == 42.0 + assert pose.quat_w[1, 0, 0].item() == -1.0 + plan.close() + + +def test_scene_plan_validates_body_field_requests_and_pose_reads( + monkeypatch: pytest.MonkeyPatch, +) -> None: + backend = _SceneTensorBackend() + _use_body_capable_tensor_spec(monkeypatch) + scene = _scene(backend) + + with pytest.raises(ValueError, match="body_fields are unknown"): + SceneTensorReadSpec(entity="robot", body_names=("hip",), body_fields=("pos_w", "bogus")) + with pytest.raises(ValueError, match="body_fields must be unique"): + SceneTensorReadSpec(entity="robot", body_names=("hip",), body_fields=("pos_w", "pos_w")) + with pytest.raises(ValueError, match="body_fields cannot be empty"): + SceneTensorReadSpec(entity="robot", body_names=("hip",), body_fields=()) + + plan = scene.compile_tensor_reads( + "cpu", + (SceneTensorReadSpec(entity="robot", body_names=("hip",), body_fields=("pos_w",)),), + ) + with pytest.raises(RuntimeError, match="must be refreshed"): + plan.ray_pose_view("robot") + plan.refresh() + with pytest.raises(ValueError, match=r"fields \['quat_w'\] were not compiled"): + plan.ray_pose_view("robot") + plan.close() + + def test_scene_plan_rejects_reads_before_refresh_and_after_close() -> None: backend = _SceneTensorBackend() plan = _scene(backend).compile_tensor_reads("cpu", _specs()) diff --git a/tests/envs/mdp/test_ray_observation.py b/tests/envs/mdp/test_ray_observation.py new file mode 100644 index 000000000..8330aa51c --- /dev/null +++ b/tests/envs/mdp/test_ray_observation.py @@ -0,0 +1,278 @@ +"""Focused tests for the tensor-native ray observation integration.""" + +from __future__ import annotations + +import sys +from dataclasses import dataclass +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest +import torch + +from unilab.envs.mdp.ray_observation import RayPatternCfg, RayScan +from unilab.managers.manager_base import ManagerTermBaseCfg + +_REPO_ROOT = Path(__file__).resolve().parents[3] + + +class _Block: + """Minimal RayTraceTensorResult-shaped caster output.""" + + def __init__(self, values: torch.Tensor): + self.distance = values + + +@dataclass +class _Capabilities: + supports_tensor_pose_sync: bool = True + supports_tensor_rays: bool = True + supports_device_output: bool = True + stream_event_ownership: str | None = "caller" + torch_devices: tuple[str, ...] = ("cpu",) + + +class _FakeCaster: + caster_type = "fake-ray" + + def __init__(self, *, values: torch.Tensor | None = None): + self.values = values + self.scene: Any = None + self.materialized = False + self.layout: tuple[str, ...] | None = None + self.poses: list[Any] = [] + self.traces = 0 + self.closed = 0 + + @property + def num_envs(self) -> int: + return 2 + + @property + def num_rays(self) -> int: + return 4 + + def get_ray_capabilities(self) -> _Capabilities: + return _Capabilities() + + def materialize(self, scene: Any) -> None: + assert scene is self.scene + self.materialized = True + + def bind_pose_layout(self, body_names) -> None: + self.layout = tuple(body_names) + + def update_pose_tensor(self, pose: Any) -> None: + self.poses.append(pose) + + def trace_tensor(self, origins, directions, max_distance, env_ids=None, outputs=None): + self.traces += 1 + self.last_inputs = (origins, directions, max_distance) + values = self.values + if values is None: + values = torch.arange(8, dtype=torch.float32).reshape(2, 4) + return _Block(values) + + def close(self) -> None: + self.closed += 1 + + +class _RayScene: + num_bodies = 1 + num_geoms = 0 + geom_types = () + geom_sizes = np.empty((0, 3), dtype=np.float64) + geom_local_pos = np.empty((0, 3), dtype=np.float64) + geom_local_quat = np.empty((0, 4), dtype=np.float64) + geom_body_ids = np.empty((0,), dtype=np.intp) + body_names = ("base",) + + +class _ReadPlan: + device = torch.device("cpu") + ready = True + + def __init__(self, pos: torch.Tensor, quat: torch.Tensor): + self.pos_w = pos + self.quat_w = quat + + def ray_pose_view(self, entity_name: str, body_names) -> Any: + self.entity_name = entity_name + self.body_names = tuple(body_names) + return SimpleNamespace(body_names=self.body_names, pos_w=self.pos_w, quat_w=self.quat_w) + + +class _Backend: + backend_type = "fake-ray-backend" + num_envs = 2 + + def __init__(self, scene: Any | None = None): + self.scene = scene if scene is not None else _RayScene() + + def get_ray_scene_description(self, body_names=None) -> Any: + if tuple(body_names or ()) != self.scene.body_names: + raise KeyError(tuple(body_names)) + return self.scene + + +class _Env: + num_envs = 2 + device = torch.device("cpu") + backend = _Backend() + + def __init__(self, read_plan: Any, caster: _FakeCaster, scene: Any): + self.scene = SimpleNamespace(read_plan=read_plan) + self.backend = _Backend(scene) + self._caster = caster + + @property + def caster(self) -> _FakeCaster: + return self._caster + + +def _bind(monkeypatch, values: torch.Tensor | None = None) -> tuple[_FakeCaster, RayScan]: + scene = _RayScene() + caster = _FakeCaster(values=values) + read_plan = _ReadPlan( + torch.randn((2, 1, 3), dtype=torch.float32), + torch.randn((2, 1, 4), dtype=torch.float32).renorm(dim=-1, p=2, maxnorm=1.0), + ) + env = _Env(read_plan, caster, scene) + pattern = RayPatternCfg( + entity_name="robot", + body_names=("base",), + num_rays=4, + pattern_start=0.0, + pattern_stop=1.5, + max_distance=3.0, + ) + term = RayScan(ManagerTermBaseCfg(func=RayScan, params={"ray_pattern": pattern}), env) + monkeypatch.setattr( + "unilab.envs.mdp.ray_observation._create_ray_caster", + lambda caster_name, **kwargs: (setattr(caster, "scene", scene), caster)[1], + ) + term.bind_read_phase() + return caster, term + + +def test_ray_pattern_config_validates_public_names_and_values() -> None: + assert RayPatternCfg( + entity_name="robot", + body_names=("base",), + num_rays=4, + pattern_start=0, + pattern_stop=1, + max_distance=2, + ).body_names == ("base",) + for kwargs in ( + {"entity_name": ""}, + {"body_names": "base"}, + {"body_names": ()}, + {"body_names": ("base", "base")}, + {"num_rays": 0}, + {"max_distance": 0.0}, + {"caster": ""}, + ): + base = dict( + entity_name="robot", + body_names=("base",), + num_rays=4, + pattern_start=0.0, + pattern_stop=1.0, + max_distance=1.0, + ) + base.update(kwargs) + with pytest.raises((TypeError, ValueError)): + RayPatternCfg(**base) + + +def test_ray_scan_binds_and_traces_without_importing_uni_ray(monkeypatch) -> None: + class _Blocker: + def find_module(self, fullname, path=None): + if fullname == "uni_ray" or fullname.startswith("uni_ray."): + return self + return None + + def load_module(self, fullname): + raise AssertionError(f"uni_ray was imported: {fullname}") + + sys.meta_path.insert(0, _Blocker()) + try: + caster, term = _bind(monkeypatch) + result = term(term._env) + + assert caster.materialized + assert caster.layout == ("base",) + assert term.tensor_body_names == ("base",) + assert term.observation_width == 4 + assert caster.poses[-1].body_names == ("base",) + assert result.shape == (2, 4) + assert result.dtype == torch.float32 + assert result.device == torch.device("cpu") + term.close() + assert caster.closed == 1 + term.close() + assert caster.closed == 1 + finally: + del sys.meta_path[0] + + +def test_ray_scan_copies_reused_caster_output_into_term_owned_carrier(monkeypatch) -> None: + source = torch.arange(8, dtype=torch.float32).reshape(2, 4) + caster, term = _bind(monkeypatch, source) + + first = term(term._env) + second = term(term._env) + + assert first is second + torch.testing.assert_close(first, source) + source[0, 0] = 99.0 + third = term(term._env) + assert third[0, 0].item() == source[0, 0].item() + term.close() + + +def test_ray_scan_fails_closed_before_binding_and_on_unsupported_capability(monkeypatch) -> None: + scene = _RayScene() + caster = _FakeCaster() + read_plan = _ReadPlan( + torch.zeros((2, 1, 3), dtype=torch.float32), + torch.zeros((2, 1, 4), dtype=torch.float32), + ) + env = _Env(read_plan, caster, scene) + pattern = RayPatternCfg( + entity_name="robot", + body_names=("base",), + num_rays=2, + pattern_start=0.0, + pattern_stop=1.0, + max_distance=1.0, + ) + term = RayScan(ManagerTermBaseCfg(func=RayScan, params={"ray_pattern": pattern}), env) + with pytest.raises(NotImplementedError, match="bind its read phase"): + term(env) + + caster.get_ray_capabilities = lambda: _Capabilities(supports_tensor_rays=False) + monkeypatch.setattr( + "unilab.envs.mdp.ray_observation._create_ray_caster", + lambda caster_name, **kwargs: caster, + ) + with pytest.raises(NotImplementedError, match="tensor pose/ray/output profile"): + term.bind_read_phase() + assert not caster.materialized + assert caster.layout is None + + +def test_ray_scan_uses_only_public_backend_and_factory_boundaries(monkeypatch) -> None: + caster, term = _bind(monkeypatch) + env = term._env + + scene = env.backend.get_ray_scene_description(("base",)) + assert scene is env.backend.scene + assert term.ray_pattern.caster == "uni_ray" + assert not any( + path.name.startswith("_") and path.is_file() for path in Path(__file__).parent.rglob("*.py") + ) + term.close() diff --git a/tests/managers/test_device_contracts.py b/tests/managers/test_device_contracts.py index dfdde5429..f8c760ead 100644 --- a/tests/managers/test_device_contracts.py +++ b/tests/managers/test_device_contracts.py @@ -16,6 +16,7 @@ TensorExecution, TensorLifecycleCapabilities, TensorProcessTopology, + TrackedBodyStateViews, ) from unilab.base.entity import EntityCfg, EntityScene, SceneTensorReadSpec @@ -308,6 +309,18 @@ def get_state_views(self, names, device=None): def get_sensor_view(self, name, device=None): return torch.as_tensor(self.sensor_values[name], dtype=torch.float32, device=device) + def get_tracked_body_views(self, body_names, device=None): + names = tuple(name.rsplit("/", 1)[-1] for name in body_names) + if names != ("base",): + raise KeyError(names) + return TrackedBodyStateViews( + body_names=names, + pos_w=torch.as_tensor(self.body_pos, dtype=torch.float32, device=device), + quat_w=torch.as_tensor(self.body_quat, dtype=torch.float32, device=device), + lin_vel_w=torch.as_tensor(self.body_lin_vel_w, dtype=torch.float32, device=device), + ang_vel_w=torch.as_tensor(self.body_ang_vel_w, dtype=torch.float32, device=device), + ) + def get_joint_dof_pos_indices(self, names) -> np.ndarray: return np.asarray([self.joint_names.index(name) for name in names], dtype=np.int32) diff --git a/tests/scripts/test_rocm_dependency_profile.py b/tests/scripts/test_rocm_dependency_profile.py index bfc1a7c27..ed954c2d8 100644 --- a/tests/scripts/test_rocm_dependency_profile.py +++ b/tests/scripts/test_rocm_dependency_profile.py @@ -24,6 +24,10 @@ def _project_scripts(path: Path) -> dict[str, str]: return dict(project["project"]["scripts"]) +def _project(path: Path) -> dict: + return tomllib.loads(path.read_text(encoding="utf-8")) + + def test_rocm_profile_preserves_non_substituted_runtime_dependencies() -> None: default_dependencies = _dependencies(ROOT / "pyproject.toml") rocm_dependencies = _dependencies(ROOT / "pyproject.rocm.toml") @@ -55,3 +59,19 @@ def test_rocm_lock_contains_profile_runtime_dependencies() -> None: assert "prettytable" in packages assert "prettytable" in root_dependencies assert "numba" not in root_dependencies + + +def test_ray_local_optional_dependencies_stay_branch_local() -> None: + for path in (ROOT / "pyproject.toml", ROOT / "pyproject.rocm.toml"): + project = _project(path) + extras = project["project"]["optional-dependencies"] + assert extras["ray-local"] == ["uni-ray"] + assert "uni-ray" not in _dependencies(path) + + sources = project["tool"]["uv"]["sources"] + assert sources["unisim-core"] == [ + {"path": "../unisim", "editable": True, "marker": "extra == 'ray-local'"} + ] + assert sources["uni-ray"] == [ + {"path": "../uni_ray", "editable": True, "marker": "extra == 'ray-local'"} + ] From 4024743296274d8665706d169bc1263f018c3b3d Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Mon, 26 Oct 2026 03:24:53 +0800 Subject: [PATCH 2/4] fix: publish shared ray batches with matching shapes --- src/unilab/envs/mdp/ray_observation.py | 5 ++-- tests/envs/mdp/test_ray_observation.py | 34 ++++++++++++++++++++++++++ 2 files changed, 36 insertions(+), 3 deletions(-) diff --git a/src/unilab/envs/mdp/ray_observation.py b/src/unilab/envs/mdp/ray_observation.py index 27d39c5c9..88e26298e 100644 --- a/src/unilab/envs/mdp/ray_observation.py +++ b/src/unilab/envs/mdp/ray_observation.py @@ -142,12 +142,11 @@ def bind_read_phase(self) -> None: (torch.cos(angles), torch.sin(angles), torch.zeros_like(angles)), dim=-1 ).contiguous() self._origins = torch.zeros( - (self._env.num_envs, self.ray_pattern.num_rays, 3), + (self.ray_pattern.num_rays, 3), dtype=torch.float32, device=self._device, ) - self._origins[:, :, 2] = self.ray_pattern.pattern_center_height - self._origins = self._origins.contiguous() + self._origins[:, 2] = self.ray_pattern.pattern_center_height self._bound = True def __call__(self, env: ManagerBasedRlEnv) -> torch.Tensor: diff --git a/tests/envs/mdp/test_ray_observation.py b/tests/envs/mdp/test_ray_observation.py index 8330aa51c..14ee478fc 100644 --- a/tests/envs/mdp/test_ray_observation.py +++ b/tests/envs/mdp/test_ray_observation.py @@ -70,6 +70,19 @@ def update_pose_tensor(self, pose: Any) -> None: def trace_tensor(self, origins, directions, max_distance, env_ids=None, outputs=None): self.traces += 1 self.last_inputs = (origins, directions, max_distance) + if origins.shape != directions.shape: + raise ValueError( + f"origin and direction tensor batches must have the same shape: " + f"{tuple(origins.shape)} != {tuple(directions.shape)}" + ) + if origins.ndim != 2 or origins.shape[-1] != 3: + raise ValueError( + f"shared ray tensors must have shape [num_rays, 3]: {tuple(origins.shape)}" + ) + if origins.shape[0] != self.num_rays: + raise ValueError(f"expected {self.num_rays} rays, got {origins.shape[0]}") + if not origins.is_contiguous() or not directions.is_contiguous(): + raise ValueError("shared ray tensors must be contiguous") values = self.values if values is None: values = torch.arange(8, dtype=torch.float32).reshape(2, 4) @@ -234,6 +247,27 @@ def test_ray_scan_copies_reused_caster_output_into_term_owned_carrier(monkeypatc term.close() +def test_ray_scan_publishes_matching_shared_ray_batches(monkeypatch) -> None: + caster, term = _bind(monkeypatch) + + term(term._env) + + origins, directions, max_distance = caster.last_inputs + assert origins.shape == directions.shape == (4, 3) + assert origins.is_contiguous() and directions.is_contiguous() + assert origins.dtype == directions.dtype == torch.float32 + assert origins.device == directions.device == torch.device("cpu") + torch.testing.assert_close( + origins, + torch.tensor( + [[0.0, 0.0, term.ray_pattern.pattern_center_height]] * 4, + dtype=torch.float32, + ), + ) + assert max_distance == term.ray_pattern.max_distance + term.close() + + def test_ray_scan_fails_closed_before_binding_and_on_unsupported_capability(monkeypatch) -> None: scene = _RayScene() caster = _FakeCaster() From f524f4aedc6738fa2e79d7040c49eac7dd22ff95 Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Mon, 26 Oct 2026 03:28:20 +0800 Subject: [PATCH 3/4] fix: bind packed ray pose views to the public host packet --- src/unilab/base/entity.py | 34 ++++++++++++++++++++-- tests/base/test_entity_tensor_read_plan.py | 10 +++++-- 2 files changed, 40 insertions(+), 4 deletions(-) diff --git a/src/unilab/base/entity.py b/src/unilab/base/entity.py index 54c163243..736a28271 100644 --- a/src/unilab/base/entity.py +++ b/src/unilab/base/entity.py @@ -449,7 +449,14 @@ def refresh(self) -> None: """Publish one full-batch packet for the current read phase.""" self._require_not_closed() if self._host_plan_value is not None: - packet = self._read_packet(self._host_plan_value.read_state_sensors) + if self._packet_body_blocks: + + def reader() -> dict[str, torch.Tensor]: + return self._read_host_body_blocks(selected=False) + + else: + reader = self._host_plan_value.read_state_sensors + packet = self._read_packet(reader) else: packet = self._read_device_resident_packet() self._publish_packet(packet) @@ -463,7 +470,14 @@ def refresh_selected(self) -> None: """ self._require_not_closed() if self._host_plan_value is not None: - packet = self._read_packet(self._host_plan_value.read_selected_state_sensors) + if self._packet_body_blocks: + + def reader() -> dict[str, torch.Tensor]: + return self._read_host_body_blocks(selected=True) + + else: + reader = self._host_plan_value.read_selected_state_sensors + packet = self._read_packet(reader) else: packet = self._read_device_resident_packet() self._publish_packet(packet) @@ -479,6 +493,22 @@ def _read_device_resident_packet(self) -> dict[str, torch.Tensor]: """ return self._read_packet(None) + def _read_host_body_blocks(self, *, selected: bool) -> dict[str, torch.Tensor]: + """Expose packed body blocks as views of the single host-bridge packet.""" + host_plan = self._host_plan_value + assert host_plan is not None + source = ( + host_plan.read_selected_state_sensors() if selected else host_plan.read_state_sensors() + ) + packet = dict(source) + poses = host_plan.get_body_pose_views() + for field in _TENSOR_POSE_BODY_FIELDS: + packet_name = _TENSOR_BODY_PACKET_FIELDS[field] + value = getattr(poses, field) + if value is not None: + packet[packet_name] = value.reshape(value.shape[0], -1) + return packet + def invalidate(self) -> None: """Drop phase values after an in-phase simulation mutation.""" self._require_not_closed() diff --git a/tests/base/test_entity_tensor_read_plan.py b/tests/base/test_entity_tensor_read_plan.py index 01f1b2b6f..4dd89f3e9 100644 --- a/tests/base/test_entity_tensor_read_plan.py +++ b/tests/base/test_entity_tensor_read_plan.py @@ -3,7 +3,7 @@ from __future__ import annotations from dataclasses import dataclass, field -from types import MethodType +from types import MethodType, SimpleNamespace from typing import cast import numpy as np @@ -219,6 +219,13 @@ def read_selected_state_sensors(self): backend.selected_reads += 1 return backend._packet() + def get_body_pose_views(self): + return SimpleNamespace( + body_names=("robot/hip",), + pos_w=backend.body_blocks["body_pos_w"], + quat_w=backend.body_blocks["body_quat_w"], + ) + def write_control(self, ctrl) -> None: return None @@ -239,7 +246,6 @@ def _packet(self) -> dict[str, torch.Tensor]: packet.update(self.sensors) if self.device_resident: return packet - packet.update(self.body_blocks) return packet From b825ca8658030a4c816b96d9cf6e7f7256eb127b Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Mon, 26 Oct 2026 16:36:31 +0800 Subject: [PATCH 4/4] feat: integrate public rich ray perception queries --- docs/sphinx/source/changelog.md | 44 ++ .../source/en/5-reference/5-support_matrix.md | 2 +- .../zh_CN/5-reference/5-support_matrix.md | 2 +- pyproject.rocm.toml | 8 +- pyproject.toml | 10 +- src/unilab/base/entity.py | 49 ++- src/unilab/envs/manager_based_rl_env.py | 23 +- src/unilab/envs/mdp/ray_observation.py | 371 ++++++++++++++-- tests/base/test_entity_tensor_read_plan.py | 21 +- tests/envs/mdp/test_ray_observation.py | 409 ++++++++++++++++-- tests/envs/test_manager_based_rl_env.py | 87 ++++ tests/scripts/test_rocm_dependency_profile.py | 2 +- uv.lock | 93 +++- uv.rocm.lock | 105 ++++- 14 files changed, 1077 insertions(+), 149 deletions(-) diff --git a/docs/sphinx/source/changelog.md b/docs/sphinx/source/changelog.md index 9e3c2bdfe..096cdf227 100644 --- a/docs/sphinx/source/changelog.md +++ b/docs/sphinx/source/changelog.md @@ -13,6 +13,50 @@ UniLab 遵循[语义化版本](https://semver.org/)。本共享页面以中英 PyPI 版本变更与未发布变更；发布日期采用 PyPI 上传日期。完整提交历史请参阅 [UniLab 仓库](https://github.com/unilabsim/UniLab)。 +## Unreleased / 未发布 + +### Added / 新增 + +- The Manager tensor ray observation now consumes UniSim's full-scene public + `RayCollisionDescription`, so mesh/hfield materialization no longer requires + a backend-specific `uni_ray` builder, and entity-local packed pose reads are + aggregated onto the public full-scene body axis. `RayPatternCfg` supports a + body-only attachment transform with a local sensor offset, `min_distance`, + public geom/tag/body/ancestor filters, a yaw ring with configurable pitch, + and optional distance+normal output + (`[distance, normal_x, normal_y, normal_z]`, flattened to the observation + width). The branch-local `ray-local` + dependency profile resolves sibling `unisim-core` and `uni-ray` through + relative editable paths in both CUDA and ROCm lockfiles; no PyPI publication + is involved + ([UniSim #472](https://github.com/unilabsim/unisim/issues/472)). + Manager tensor ray observation 现在消费 UniSim 公开 full-scene + `RayCollisionDescription`，mesh/hfield materialization 不再要求 backend-specific + `uni_ray` builder，并将 entity-local packed pose read 聚合到公开 full-scene body + 轴。`RayPatternCfg` 支持 body-only attachment 变换与本地 sensor + offset、`min_distance`、公开 geom/tag/body/ancestor filter，以及可选 + distance+normal 输出（`[distance, normal_x, normal_y, normal_z]`）。branch-local + `ray-local` 依赖 profile 在 CUDA 与 ROCm lockfile 中均通过相对 editable path + 解析 sibling `unisim-core` 与 `uni-ray`；不涉及 PyPI 发布 + （[UniSim #472](https://github.com/unilabsim/unisim/issues/472)）。 + The first facade is a minimal grouped ray sensor, not a camera: intrinsics, + image projection, noise, history, and scheduling remain application logic. + Full-scene pose aggregation covers entity-owned bodies plus the identity + world root; uncovered non-root bodies fail closed. MuJoCo and MJWarp + fixed-variant collision scenes and Motrix source mesh/hfield scenes follow + the UniSim fail-closed boundaries. The real end-to-end fixture covers + primitive geometry, static filtering, minimum distance, attachment + transformation, and normals; mesh/hfield numerical parity is owned by the + UniSim and `uni_ray` contract suites. + 首个 facade 是最小 grouped ray sensor，不是 camera：intrinsics、图像投影、 + noise、history 与 scheduling 仍属于应用逻辑。full-scene pose 聚合覆盖 + entity-owned body 与 identity world root；未覆盖的非 root body fail closed。 + MuJoCo 与 MJWarp fixed-variant collision scene 以及 Motrix source + mesh/hfield scene 遵循 UniSim fail-closed 边界。真实端到端 fixture 覆盖 + primitive geometry、static filtering、minimum distance、attachment + transformation 与 normal；mesh/hfield 数值 parity 由 UniSim 与 `uni_ray` + contract suite 负责。 + ## 1.5.1 (2026-10-23) ### Added / 新增 diff --git a/docs/sphinx/source/en/5-reference/5-support_matrix.md b/docs/sphinx/source/en/5-reference/5-support_matrix.md index e0fcfef3a..dec1400c3 100644 --- a/docs/sphinx/source/en/5-reference/5-support_matrix.md +++ b/docs/sphinx/source/en/5-reference/5-support_matrix.md @@ -81,7 +81,7 @@ This table is derived from UniSim's SDK-free public static inventory. It describ | `mjwarp` | Device-resident / in-process / direct | CUDA | Required for the entire tensor lifecycle | Supported: Linux CUDA only | Unsupported; no CPU, MPS, or ROCm fallback | Unsupported; no CPU, MPS, or ROCm fallback | In-process; no external Python worker | Unsupported | Unsupported | Unsupported | Unsupported | | `newton` | Device-resident / in-process / direct | CUDA | Required for the entire tensor lifecycle | Supported: Linux CUDA only | Unsupported; no CPU, MPS, or ROCm fallback | Unsupported; no CPU, MPS, or ROCm fallback | In-process; no external Python worker | Exact | Unsupported | Unsupported | Unsupported | | `superdex` | Host bridge / in-process / host bridge | CPU / CUDA | Required only when the learner requests CUDA state/control buffers | Supported: CPU-authoritative physics with optional CUDA Torch buffers | CPU-authoritative host bridge only; no CUDA physics claim | CPU-authoritative host bridge with explicit ROCm Torch buffers (`cuda`) | In-process; no external Python worker | Unsupported | Unsupported | Unsupported | Exact | -| `genesis` | Device-resident / in-process / direct | CUDA | Required for the entire tensor lifecycle | Supported: Linux CUDA only | Unsupported; no CPU, MPS, or ROCm fallback | Unsupported; no CPU, MPS, or ROCm fallback | In-process; no external Python worker | Exact | Unsupported | Unsupported | Unsupported | +| `genesis` | Device-resident / in-process / direct | CUDA | Required for the entire tensor lifecycle | Supported: Linux CUDA only | Unsupported UniSim tensor profile; Genesis engine natively supports Metal, with no CPU or substitute-backend fallback | Unsupported UniSim tensor profile; Genesis engine natively supports AMDGPU, with no CPU or substitute-backend fallback | In-process; no external Python worker | Exact | Unsupported | Unsupported | Unsupported | | `isaacgym` | Device-resident / external worker / cuda ipc | CUDA | Required for the entire tensor lifecycle | Supported: Linux CUDA only | Unsupported; no CPU, MPS, or ROCm fallback | Unsupported; no CPU, MPS, or ROCm fallback | Dedicated external Python 3.8 worker; host Python paths are not inherited | Exact | Unsupported | Unsupported | Unsupported | | `isaacsim` | Device-resident / external worker / cuda ipc | CUDA | Required for the entire tensor lifecycle | Supported: Linux CUDA only | Unsupported; no CPU, MPS, or ROCm fallback | Unsupported; no CPU, MPS, or ROCm fallback | Dedicated external Python 3.11 worker; host Python paths are not inherited | Exact | Unsupported | Unsupported | Unsupported | diff --git a/docs/sphinx/source/zh_CN/5-reference/5-support_matrix.md b/docs/sphinx/source/zh_CN/5-reference/5-support_matrix.md index fdb0f83ac..e6cd0c6d5 100644 --- a/docs/sphinx/source/zh_CN/5-reference/5-support_matrix.md +++ b/docs/sphinx/source/zh_CN/5-reference/5-support_matrix.md @@ -73,7 +73,7 @@ uv run scripts/generate_support_matrix.py --write | `mjwarp` | Device-resident / in-process / direct | CUDA | Required for the entire tensor lifecycle | Supported: Linux CUDA only | Unsupported; no CPU, MPS, or ROCm fallback | Unsupported; no CPU, MPS, or ROCm fallback | In-process; no external Python worker | 不支持 | 不支持 | 不支持 | 不支持 | | `newton` | Device-resident / in-process / direct | CUDA | Required for the entire tensor lifecycle | Supported: Linux CUDA only | Unsupported; no CPU, MPS, or ROCm fallback | Unsupported; no CPU, MPS, or ROCm fallback | In-process; no external Python worker | 支持 | 不支持 | 不支持 | 不支持 | | `superdex` | Host bridge / in-process / host bridge | CPU / CUDA | Required only when the learner requests CUDA state/control buffers | Supported: CPU-authoritative physics with optional CUDA Torch buffers | CPU-authoritative host bridge only; no CUDA physics claim | CPU-authoritative host bridge with explicit ROCm Torch buffers (`cuda`) | In-process; no external Python worker | 不支持 | 不支持 | 不支持 | 支持 | -| `genesis` | Device-resident / in-process / direct | CUDA | Required for the entire tensor lifecycle | Supported: Linux CUDA only | Unsupported; no CPU, MPS, or ROCm fallback | Unsupported; no CPU, MPS, or ROCm fallback | In-process; no external Python worker | 支持 | 不支持 | 不支持 | 不支持 | +| `genesis` | Device-resident / in-process / direct | CUDA | Required for the entire tensor lifecycle | Supported: Linux CUDA only | Unsupported UniSim tensor profile; Genesis engine natively supports Metal, with no CPU or substitute-backend fallback | Unsupported UniSim tensor profile; Genesis engine natively supports AMDGPU, with no CPU or substitute-backend fallback | In-process; no external Python worker | 支持 | 不支持 | 不支持 | 不支持 | | `isaacgym` | Device-resident / external worker / cuda ipc | CUDA | Required for the entire tensor lifecycle | Supported: Linux CUDA only | Unsupported; no CPU, MPS, or ROCm fallback | Unsupported; no CPU, MPS, or ROCm fallback | Dedicated external Python 3.8 worker; host Python paths are not inherited | 支持 | 不支持 | 不支持 | 不支持 | | `isaacsim` | Device-resident / external worker / cuda ipc | CUDA | Required for the entire tensor lifecycle | Supported: Linux CUDA only | Unsupported; no CPU, MPS, or ROCm fallback | Unsupported; no CPU, MPS, or ROCm fallback | Dedicated external Python 3.11 worker; host Python paths are not inherited | 支持 | 不支持 | 不支持 | 不支持 | diff --git a/pyproject.rocm.toml b/pyproject.rocm.toml index f10e8c29a..f1deb004a 100644 --- a/pyproject.rocm.toml +++ b/pyproject.rocm.toml @@ -23,7 +23,7 @@ classifiers = [ requires-python = ">=3.10,<3.14" dependencies = [ "numpy", - "unisim-core==2.1.1", + "unisim-core>=2.1.1", "prettytable>=3.10", "torch==2.14.0", "triton-rocm==3.8.0 ; sys_platform == 'linux' and platform_machine == 'x86_64'", @@ -72,11 +72,13 @@ mujoco = [ "mujoco~=3.11.0", "mjbatch-uni~=0.2.5", ] -motrix = ["unisim-core[motrix]==2.1.1"] +motrix = ["unisim-core[motrix]>=2.1.1"] # Branch-local ROCm integration against sibling ray-contract checkouts. Keep # this extra aligned with pyproject.toml; the default runtime still resolves the # published unisim-core pin. uni-ray never becomes a runtime dependency. -ray-local = ["uni-ray"] +# On this issue branch the generated lock intentionally resolves sibling editable +# checkouts; a standard checkout must provide ../unisim and ../uni_ray. +ray-local = ["uni-ray[warp]"] [dependency-groups] # The RL stack (APPO/off-policy/multi-GPU) is merged into unilab itself diff --git a/pyproject.toml b/pyproject.toml index 5a5440364..cc3d97295 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,7 +33,7 @@ classifiers = [ requires-python = ">=3.10,<3.14" dependencies = [ "numpy", - "unisim-core==2.1.1", + "unisim-core>=2.1.1", "prettytable>=3.10", # A range (not an exact pin) lets ROCm users substitute a ROCm torch # build; uv.lock pins the tested CUDA builds via tool.uv.sources. @@ -106,18 +106,20 @@ newton = [ "pyglet>=2.1.6,<3", "imgui-bundle>=1.92.0", ] -motrix = ["unisim-core[motrix]==2.1.1"] +motrix = ["unisim-core[motrix]>=2.1.1"] genesis = ["genesis-world==1.3.3"] # Wheels are CPython 3.12/3.13 Linux x86_64 only; the marker keeps uv lock # resolvable for the other required-environments. superdex = [ - "unisim-core[superdex]==2.1.1 ; python_version >= '3.12' and sys_platform == 'linux' and platform_machine == 'x86_64'", + "unisim-core[superdex]>=2.1.1 ; python_version >= '3.12' and sys_platform == 'linux' and platform_machine == 'x86_64'", ] # Branch-local integration against sibling ray-contract checkouts. The exact # published UniSim pin remains the default; the local path source overrides it # only when this optional extra is selected. uni-ray never becomes a runtime # dependency and is imported only by UniSim's public caster factory. -ray-local = ["uni-ray"] +# On this issue branch the generated lock intentionally resolves sibling editable +# checkouts; a standard checkout must provide ../unisim and ../uni_ray. +ray-local = ["uni-ray[warp]"] [dependency-groups] # The RL stack (APPO/off-policy/multi-GPU) is merged into unilab itself diff --git a/src/unilab/base/entity.py b/src/unilab/base/entity.py index 736a28271..a849d8df5 100644 --- a/src/unilab/base/entity.py +++ b/src/unilab/base/entity.py @@ -455,7 +455,10 @@ def reader() -> dict[str, torch.Tensor]: return self._read_host_body_blocks(selected=False) else: - reader = self._host_plan_value.read_state_sensors + reader = cast( + Callable[[], dict[str, torch.Tensor]], + self._host_plan_value.read_state_sensors, + ) packet = self._read_packet(reader) else: packet = self._read_device_resident_packet() @@ -476,7 +479,10 @@ def reader() -> dict[str, torch.Tensor]: return self._read_host_body_blocks(selected=True) else: - reader = self._host_plan_value.read_selected_state_sensors + reader = cast( + Callable[[], dict[str, torch.Tensor]], + self._host_plan_value.read_selected_state_sensors, + ) packet = self._read_packet(reader) else: packet = self._read_device_resident_packet() @@ -810,8 +816,12 @@ def _packet_body_fields( # ``reshape`` would silently materialize a contiguous copy when the # backend publishes one packed body block beside unrelated columns. value = raw.unflatten(1, (len(axis), width)) - owner_axis = tuple(_packet_body_name(owner.name, name) for name in compiled_bodies) - requested_axis = tuple(_packet_body_name(owner.name, name) for name in requested) + if self._host_plan_value is not None: + owner_axis = compiled_bodies + requested_axis = requested + else: + owner_axis = tuple(_packet_body_name(owner.name, name) for name in compiled_bodies) + requested_axis = tuple(_packet_body_name(owner.name, name) for name in requested) if requested_axis != owner_axis or owner_axis != axis: rows = self._packet_body_indices(requested_axis, axis) value = value.index_select(1, rows) @@ -4606,6 +4616,7 @@ def compile_tensor_reads( host_body_blocks = {"body_names", "body_fields"}.issubset( inspect.signature(TensorIOSpec).parameters ) + native_body_names: tuple[str, ...] = () host_pose_only_blocks = host_body_blocks and all( not (set(fields) - set(_TENSOR_POSE_BODY_FIELDS)) for fields in per_entity_body_fields.values() @@ -4614,20 +4625,18 @@ def compile_tensor_reads( packet_body_blocks = True packet_names = (*sensor_packet_names, *packet_body_names) host_sensor_packet_names = sensor_packet_names - for entity_name, body_names_value in expanded_bodies.items(): - if not body_names_value: - continue - for field in per_entity_body_fields[entity_name]: - packet_name = _TENSOR_BODY_PACKET_FIELDS[field] - axis = packet_body_axes.setdefault(packet_name, ()) - packet_body_axes[packet_name] = ( - *axis, - *( - _packet_body_name(entity_name, body_name) - for body_name in body_names_value - if _packet_body_name(entity_name, body_name) not in axis - ), - ) + # Host-bridge body blocks use the backend's native public body + # namespace. Entity ownership remains available through the + # per-entity read-plan mappings below. + native_body_names = tuple( + dict.fromkeys( + body_name + for body_names_value in expanded_bodies.values() + for body_name in body_names_value + ) + ) + for packet_name in packet_body_names: + packet_body_axes[packet_name] = native_body_names tensor_io_options: dict[str, Any] = { "state_fields": ("qpos", "qvel"), "sensor_names": ( @@ -4635,10 +4644,10 @@ def compile_tensor_reads( ), "device": resolved_device, } - if host_pose_only_blocks: + if host_pose_only_blocks and any(expanded_bodies.values()): tensor_io_options.update( { - "body_names": self._aggregate_tensor_read_body_names(expanded_bodies), + "body_names": native_body_names, "body_fields": ("pos_w", "quat_w"), } ) diff --git a/src/unilab/envs/manager_based_rl_env.py b/src/unilab/envs/manager_based_rl_env.py index 66aa466f6..8a60af6cb 100644 --- a/src/unilab/envs/manager_based_rl_env.py +++ b/src/unilab/envs/manager_based_rl_env.py @@ -541,6 +541,7 @@ def _observation_tensor_read_specs(self) -> list[SceneTensorReadSpec]: def _observation_ray_read_specs(self) -> list[SceneTensorReadSpec]: """Collect pose-only blocks needed by tensor ray observations.""" specs: list[SceneTensorReadSpec] = [] + requested_entities: set[str] = set() for group_name, terms in self.observation_manager.active_terms.items(): for name in terms: term_cfg = self.observation_manager.get_term_cfg(group_name, name) @@ -553,13 +554,23 @@ def _observation_ray_read_specs(self) -> list[SceneTensorReadSpec]: f"'{group_name}/{name}' must be RayPatternCfg; got " f"{type(ray_pattern).__name__}" ) - specs.append( - SceneTensorReadSpec( - entity=ray_pattern.entity_name, - body_names=tuple(ray_pattern.body_names), - body_fields=("pos_w", "quat_w"), + for entity_name, entity in self.scene.entities.items(): + if entity_name in requested_entities: + continue + try: + body_names = tuple(entity.body_names) + except NotImplementedError: + continue + if not body_names: + continue + requested_entities.add(entity_name) + specs.append( + SceneTensorReadSpec( + entity=entity_name, + body_names=body_names, + body_fields=("pos_w", "quat_w"), + ) ) - ) return specs def _manager_term_tensor_read_specs(self) -> list[SceneTensorReadSpec]: diff --git a/src/unilab/envs/mdp/ray_observation.py b/src/unilab/envs/mdp/ray_observation.py index 88e26298e..409823698 100644 --- a/src/unilab/envs/mdp/ray_observation.py +++ b/src/unilab/envs/mdp/ray_observation.py @@ -8,6 +8,7 @@ import torch from unisim.backend.base import torch_device_matches +from unisim.ray_query import RayPoseTensorView, RayTraceFilter, RayTraceOutputs from unilab.managers.manager_base import ManagerTermBase, ManagerTermBaseCfg @@ -20,30 +21,31 @@ class RayPatternCfg: """Declare one observation-owned ray pattern. ``caster`` is a public UniSim ray-caster factory name, not an imported - plugin class. ``body_names`` is entity-local and selects both the backend - ray-scene projection and the pose-only packed read block. + plugin class. ``attachment_body_name`` is a public full-scene body identity + from ``RayCollisionDescription.scene.body_names``; the local pattern and + sensor offset are transformed by that body pose. """ entity_name: str - body_names: Sequence[str] + attachment_body_name: str num_rays: int pattern_start: float pattern_stop: float max_distance: float + pattern_pitch: float = 0.0 caster: str = "uni_ray" - pattern_center_height: float = 0.0 + sensor_offset: Sequence[float] = (0.0, 0.0, 0.0) + min_distance: float = 0.0 + include_normal: bool = False + allowed_geom_names: Sequence[str] | None = None + allowed_tags: Sequence[str] | None = None + excluded_body_names: Sequence[str] | None = None + excluded_ancestor_body_names: Sequence[str] | None = None caster_options: dict[str, Any] = field(default_factory=dict) def __post_init__(self) -> None: if not isinstance(self.entity_name, str) or not self.entity_name: raise TypeError("RayPatternCfg entity_name must be a non-empty string") - if isinstance(self.body_names, (str, bytes)): - raise TypeError("RayPatternCfg body_names must be a sequence of strings") - names = tuple(self.body_names) - if not names or any(not isinstance(name, str) or not name for name in names): - raise TypeError("RayPatternCfg body_names must be non-empty strings") - if len(set(names)) != len(names): - raise ValueError(f"RayPatternCfg body_names must be unique: {names}") if isinstance(self.num_rays, bool) or not isinstance(self.num_rays, int): raise TypeError("RayPatternCfg num_rays must be an integer") if self.num_rays <= 0: @@ -51,8 +53,9 @@ def __post_init__(self) -> None: for name, value in ( ("pattern_start", self.pattern_start), ("pattern_stop", self.pattern_stop), + ("pattern_pitch", self.pattern_pitch), ("max_distance", self.max_distance), - ("pattern_center_height", self.pattern_center_height), + ("min_distance", self.min_distance), ): if isinstance(value, bool) or not isinstance(value, (int, float)): raise TypeError(f"RayPatternCfg {name} must be a real number") @@ -60,15 +63,55 @@ def __post_init__(self) -> None: raise ValueError(f"RayPatternCfg {name} must be finite") if self.max_distance <= 0.0: raise ValueError("RayPatternCfg max_distance must be positive") + if self.min_distance < 0.0 or self.min_distance >= self.max_distance: + raise ValueError( + "RayPatternCfg min_distance must satisfy 0 <= min_distance < max_distance" + ) + if not isinstance(self.attachment_body_name, str) or not self.attachment_body_name: + raise TypeError("RayPatternCfg attachment_body_name must be a non-empty string") + if isinstance(self.sensor_offset, (str, bytes)) or not isinstance( + self.sensor_offset, Sequence + ): + raise TypeError("RayPatternCfg sensor_offset must be a sequence of three real numbers") + offset = tuple(self.sensor_offset) + if len(offset) != 3 or any( + isinstance(value, bool) or not isinstance(value, (int, float)) for value in offset + ): + raise TypeError("RayPatternCfg sensor_offset must be a sequence of three real numbers") + if any(not math.isfinite(float(value)) for value in offset): + raise ValueError("RayPatternCfg sensor_offset must be finite") + object.__setattr__(self, "sensor_offset", tuple(float(value) for value in offset)) + if not isinstance(self.include_normal, bool): + raise TypeError("RayPatternCfg include_normal must be a boolean") if not isinstance(self.caster, str) or not self.caster: raise TypeError("RayPatternCfg caster must be a non-empty factory name") if not isinstance(self.caster_options, dict): raise TypeError("RayPatternCfg caster_options must be a mapping") - object.__setattr__(self, "body_names", names) + if "collision" in self.caster_options: + raise ValueError( + "RayPatternCfg caster_options cannot override the public collision descriptor" + ) + for name in ( + "allowed_geom_names", + "allowed_tags", + "excluded_body_names", + "excluded_ancestor_body_names", + ): + value = getattr(self, name) + if value is None: + continue + if isinstance(value, (str, bytes)) or not isinstance(value, Sequence): + raise TypeError(f"RayPatternCfg {name} must be a string sequence or None") + values = tuple(value) + if any(not isinstance(item, str) or not item for item in values): + raise TypeError(f"RayPatternCfg {name} must contain non-empty strings") + if len(set(values)) != len(values): + raise ValueError(f"RayPatternCfg {name} must be unique: {values}") + object.__setattr__(self, name, values) class RayScan(ManagerTermBase): - """Bind one tensor caster and publish its distance observation carrier.""" + """Bind one tensor caster and publish a distance/normal observation carrier.""" def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): super().__init__(env) @@ -79,17 +122,26 @@ def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv): f"got {type(pattern).__name__}" ) self.ray_pattern = pattern - self.tensor_body_names = tuple(pattern.body_names) - self.observation_width = pattern.num_rays + self.observation_width = pattern.num_rays * (4 if pattern.include_normal else 1) self._device = torch.device(env.device) self._entity_name = pattern.entity_name - self._body_names = tuple(pattern.body_names) self._caster: Any = None self._bound = False self._closed = False self._distance: torch.Tensor | None = None - self._origins: torch.Tensor | None = None - self._directions: torch.Tensor | None = None + self._normal: torch.Tensor | None = None + self._observation: torch.Tensor | None = None + self._ray_filter: RayTraceFilter | None = None + self._outputs = RayTraceOutputs() + self._attachment_body_index: int | None = None + self._local_origins: torch.Tensor | None = None + self._local_directions: torch.Tensor | None = None + self._ray_origins: torch.Tensor | None = None + self._ray_directions: torch.Tensor | None = None + self._pose_sources: dict[str, tuple[tuple[str, ...], torch.Tensor]] = {} + self._pose_body_names: tuple[str, ...] = () + self._pose_pos: torch.Tensor | None = None + self._pose_quat: torch.Tensor | None = None def bind_read_phase(self) -> None: """Materialize and bind the caster after packed pose reads compile.""" @@ -99,38 +151,79 @@ def bind_read_phase(self) -> None: if read_plan is None: raise NotImplementedError("RayScan requires the Manager packed Torch read plan") try: - scene = self._env.backend.get_ray_scene_description(self._body_names) + collision = self._env.backend.get_ray_collision_description() + scene = collision.scene + attachment_name = self.ray_pattern.attachment_body_name + if attachment_name not in scene.body_names: + raise KeyError( + f"ray pattern attachment body {attachment_name!r} is not in {scene.body_names}" + ) + self._attachment_body_index = tuple(scene.body_names).index(attachment_name) + self._pose_sources = self._resolve_pose_sources(collision) + self._validate_pose_coverage(collision) + self._pose_body_names = tuple(scene.body_names) + self._ray_filter = self._resolve_ray_filter(collision) caster = _create_ray_caster( self.ray_pattern.caster, num_envs=self._env.num_envs, num_rays=self.ray_pattern.num_rays, device=self._device, + collision=collision, **self.ray_pattern.caster_options, ) - capabilities = caster.get_ray_capabilities() - if not ( - capabilities.supports_tensor_pose_sync - and capabilities.supports_tensor_rays - and capabilities.supports_device_output - and capabilities.stream_event_ownership == "caller" - and torch_device_matches(capabilities.torch_devices, self._device) - ): - raise NotImplementedError( - f"RayScan caster '{self.ray_pattern.caster}' does not support the " - f"tensor pose/ray/output profile on {self._device}" - ) - caster.materialize(scene) - caster.bind_pose_layout(tuple(scene.body_names)) + try: + capabilities = caster.get_ray_capabilities() + if not ( + capabilities.supports_tensor_pose_sync + and capabilities.supports_tensor_rays + and capabilities.supports_device_output + and capabilities.stream_event_ownership == "caller" + and torch_device_matches(capabilities.torch_devices, self._device) + ): + raise NotImplementedError( + f"RayScan caster '{self.ray_pattern.caster}' does not support the " + f"tensor pose/ray/output profile on {self._device}" + ) + if self.ray_pattern.min_distance != 0.0 and not capabilities.supports_min_distance: + raise NotImplementedError( + f"RayScan caster '{self.ray_pattern.caster}' does not support min_distance" + ) + if self._ray_filter.has_constraints() and not capabilities.supports_static_filter: + raise NotImplementedError( + f"RayScan caster '{self.ray_pattern.caster}' does not support " + "static ray filters" + ) + if self.ray_pattern.include_normal and not capabilities.supports_tensor_normal: + raise NotImplementedError( + f"RayScan caster '{self.ray_pattern.caster}' does not support tensor normals" + ) + caster.materialize(scene) + caster.bind_pose_layout(tuple(scene.body_names)) + except Exception: + caster.close() + raise except (AttributeError, KeyError, TypeError, ValueError, NotImplementedError) as exc: raise type(exc)( f"RayScan caster '{self.ray_pattern.caster}' read-phase binding failed: {exc}" ) from exc self._caster = caster + self._outputs = RayTraceOutputs(normal=self.ray_pattern.include_normal) self._distance = torch.empty( (self._env.num_envs, self.ray_pattern.num_rays), dtype=torch.float32, device=self._device, ) + if self.ray_pattern.include_normal: + self._normal = torch.empty( + (self._env.num_envs, self.ray_pattern.num_rays, 3), + dtype=torch.float32, + device=self._device, + ) + self._observation = torch.empty( + (self._env.num_envs, self.ray_pattern.num_rays, 4), + dtype=torch.float32, + device=self._device, + ) angles = torch.linspace( self.ray_pattern.pattern_start, self.ray_pattern.pattern_stop, @@ -138,35 +231,216 @@ def bind_read_phase(self) -> None: dtype=torch.float32, device=self._device, ) - self._directions = torch.stack( - (torch.cos(angles), torch.sin(angles), torch.zeros_like(angles)), dim=-1 + pitch = float(self.ray_pattern.pattern_pitch) + self._local_directions = torch.stack( + ( + torch.cos(angles) * math.cos(pitch), + torch.sin(angles) * math.cos(pitch), + torch.full_like(angles, math.sin(pitch)), + ), + dim=-1, ).contiguous() - self._origins = torch.zeros( + self._local_origins = torch.zeros( (self.ray_pattern.num_rays, 3), dtype=torch.float32, device=self._device, ) - self._origins[:, 2] = self.ray_pattern.pattern_center_height + self._local_origins[:] = torch.tensor( + self.ray_pattern.sensor_offset, + dtype=torch.float32, + device=self._device, + ) + self._ray_origins = torch.empty( + (self._env.num_envs, self.ray_pattern.num_rays, 3), + dtype=torch.float32, + device=self._device, + ) + self._ray_directions = torch.empty_like(self._ray_origins) + self._pose_pos = torch.zeros( + (self._env.num_envs, len(scene.body_names), 3), + dtype=torch.float32, + device=self._device, + ) + self._pose_quat = torch.zeros( + (self._env.num_envs, len(scene.body_names), 4), + dtype=torch.float32, + device=self._device, + ) + self._pose_quat[..., 0] = 1.0 self._bound = True - def __call__(self, env: ManagerBasedRlEnv) -> torch.Tensor: + def __call__( + self, + env: ManagerBasedRlEnv, + ray_pattern: RayPatternCfg | None = None, + ) -> torch.Tensor: + if ray_pattern is not None and ray_pattern is not self.ray_pattern: + raise ValueError( + "RayScan was bound to a different ray_pattern; manager term parameters " + "must remain immutable" + ) if not self._bound or self._caster is None or self._distance is None: raise NotImplementedError("RayScan must bind its read phase before compute") read_plan = getattr(env.scene, "read_plan", None) if read_plan is None or not read_plan.ready: raise RuntimeError("RayScan requires a refreshed packed read phase") - pose = read_plan.ray_pose_view(self._entity_name, self._body_names) + pose = self._refresh_full_pose(read_plan) try: + self._refresh_attachment_rays(pose) self._caster.update_pose_tensor(pose) result = self._caster.trace_tensor( - self._origins, - self._directions, + self._ray_origins, + self._ray_directions, self.ray_pattern.max_distance, + outputs=self._outputs, + min_distance=self.ray_pattern.min_distance, + ray_filter=self._ray_filter, ) except (AttributeError, KeyError, TypeError, ValueError, NotImplementedError) as exc: raise type(exc)(f"RayScan tensor trace failed: {exc}") from exc - self._distance.copy_(result.distance) - return self._distance + if self._observation is None: + self._distance.copy_(result.distance) + return self._distance + assert self._normal is not None + self._observation[..., 0].copy_(result.distance) + self._normal.copy_(result.normal) + self._observation[..., 1:4].copy_(self._normal) + return self._observation.reshape(self._env.num_envs, -1) + + def _refresh_attachment_rays(self, pose: Any) -> None: + """Transform the local ray pattern by the body-only attachment pose.""" + + assert self._attachment_body_index is not None + assert self._local_origins is not None + assert self._local_directions is not None + assert self._ray_origins is not None + assert self._ray_directions is not None + index = self._attachment_body_index + position = pose.pos_w[:, index, :].unsqueeze(1) + quaternion = pose.quat_w[:, index, :] + self._ray_origins.copy_(position + _quat_rotate_wxyz(quaternion, self._local_origins)) + self._ray_directions.copy_(_quat_rotate_wxyz(quaternion, self._local_directions)) + + def _resolve_pose_sources( + self, collision: Any + ) -> dict[str, tuple[tuple[str, ...], torch.Tensor]]: + """Map public full-scene body identities to entity-local pose reads.""" + + scene = collision.scene + body_names = tuple(scene.body_names) + body_ids = self._env.backend.get_body_ids(body_names) + collision_index_by_id = { + int(body_id): index for index, body_id in enumerate(body_ids.tolist()) + } + sources: dict[str, tuple[tuple[str, ...], torch.Tensor]] = {} + entities = getattr(self._env.scene, "entities", {}) + for entity_name, entity in entities.items(): + try: + local_names = tuple(entity.body_names) + except NotImplementedError: + continue + if not local_names: + continue + local_ids = self._env.backend.get_body_ids(local_names) + selected_names: list[str] = [] + destinations: list[int] = [] + for local_name, local_id in zip(local_names, local_ids.tolist(), strict=True): + collision_index = collision_index_by_id.get(int(local_id)) + if collision_index is None: + continue + selected_names.append(local_name) + destinations.append(collision_index) + if not selected_names: + continue + sources[entity_name] = ( + tuple(selected_names), + torch.tensor(destinations, dtype=torch.long, device=self._device), + ) + return sources + + def _validate_pose_coverage(self, collision: Any) -> None: + """Require entity reads for every non-root full-scene body.""" + + scene = collision.scene + source = self._pose_sources.get(self._entity_name) + attachment_name = self.ray_pattern.attachment_body_name + if source is None or attachment_name not in source[0]: + raise KeyError( + f"ray pattern attachment body {attachment_name!r} is not owned by " + f"entity {self._entity_name!r}" + ) + covered = { + body_name for local_names, _ in self._pose_sources.values() for body_name in local_names + } + uncovered = set(scene.body_names) - covered + root_name = scene.body_names[0] + if uncovered and ( + uncovered != {root_name} + or (collision.body_parent_ids is not None and int(collision.body_parent_ids[0]) != -1) + ): + raise NotImplementedError( + "RayScan full-scene pose aggregation requires every non-root body to " + "belong to a scene entity with a packed body read; uncovered bodies: " + f"{sorted(uncovered)}" + ) + + def _refresh_full_pose(self, read_plan: Any) -> RayPoseTensorView: + """Assemble the full public body-axis pose from entity-local read views.""" + + assert self._pose_body_names + assert self._pose_pos is not None + assert self._pose_quat is not None + self._pose_pos.zero_() + self._pose_quat.zero_() + self._pose_quat[..., 0] = 1.0 + for entity_name, (local_names, destinations) in self._pose_sources.items(): + source = read_plan.ray_pose_view(entity_name, local_names) + self._pose_pos[:, destinations, :] = source.pos_w + self._pose_quat[:, destinations, :] = source.quat_w + return RayPoseTensorView( + body_names=self._pose_body_names, + pos_w=self._pose_pos, + quat_w=self._pose_quat, + ) + + def _resolve_ray_filter(self, collision: Any) -> RayTraceFilter: + """Resolve public names/tags to the compact static filter on the cold path.""" + + scene = collision.scene + geom_names = tuple(collision.geom_names) + body_names = tuple(scene.body_names) + allowed_geom_ids = None + if self.ray_pattern.allowed_geom_names is not None: + unknown = set(self.ray_pattern.allowed_geom_names) - set(geom_names) + if unknown: + raise KeyError(f"ray pattern geom names are missing: {sorted(unknown)}") + allowed_geom_ids = tuple( + geom_names.index(name) for name in self.ray_pattern.allowed_geom_names + ) + excluded_body_ids = None + if self.ray_pattern.excluded_body_names is not None: + unknown = set(self.ray_pattern.excluded_body_names) - set(body_names) + if unknown: + raise KeyError(f"ray pattern body names are missing: {sorted(unknown)}") + excluded_body_ids = tuple( + body_names.index(name) for name in self.ray_pattern.excluded_body_names + ) + excluded_ancestor_body_ids = None + if self.ray_pattern.excluded_ancestor_body_names is not None: + unknown = set(self.ray_pattern.excluded_ancestor_body_names) - set(body_names) + if unknown: + raise KeyError(f"ray pattern ancestor body names are missing: {sorted(unknown)}") + excluded_ancestor_body_ids = tuple( + body_names.index(name) for name in self.ray_pattern.excluded_ancestor_body_names + ) + ray_filter = RayTraceFilter( + allowed_geom_ids=allowed_geom_ids, + allowed_tags=self.ray_pattern.allowed_tags, + excluded_body_ids=excluded_body_ids, + excluded_ancestor_body_ids=excluded_ancestor_body_ids, + ) + ray_filter.geom_mask(collision) + return ray_filter def close(self) -> None: """Release caster resources before the scene read plan closes.""" @@ -183,6 +457,15 @@ def reset(self, env_ids: torch.Tensor | slice | None = None) -> None: __all__ = ["RayPatternCfg", "RayScan"] +def _quat_rotate_wxyz(quaternion: torch.Tensor, vector: torch.Tensor) -> torch.Tensor: + """Rotate shared local vectors by full-batch public ``wxyz`` quaternions.""" + + qvec = quaternion[:, None, 1:] + scalar = quaternion[:, None, 0:1] + cross = 2.0 * torch.linalg.cross(qvec, vector[None, :, :]) + return vector[None, :, :] + scalar * cross + torch.linalg.cross(qvec, cross) + + def _create_ray_caster(caster: str, **options: Any) -> Any: """Lazily call the installed UniSim public caster factory.""" import unisim diff --git a/tests/base/test_entity_tensor_read_plan.py b/tests/base/test_entity_tensor_read_plan.py index 4dd89f3e9..79b0ddce3 100644 --- a/tests/base/test_entity_tensor_read_plan.py +++ b/tests/base/test_entity_tensor_read_plan.py @@ -274,6 +274,18 @@ class _BodyCapableTensorIOSpec(TensorIOSpec): monkeypatch.setattr("unilab.base.entity.TensorIOSpec", _BodyCapableTensorIOSpec) +def _use_legacy_tensor_spec(monkeypatch: pytest.MonkeyPatch) -> None: + """Expose a pre-2.2 TensorIOSpec without packed body-block parameters.""" + + @dataclass(frozen=True) + class _LegacyTensorIOSpec: + state_fields: tuple[str, ...] + sensor_names: tuple[str, ...] + device: object | None = None + + monkeypatch.setattr("unilab.base.entity.TensorIOSpec", _LegacyTensorIOSpec) + + def _specs() -> tuple[SceneTensorReadSpec, ...]: return ( SceneTensorReadSpec(entity="robot", sensor_names=("imu_gyro",)), @@ -313,9 +325,12 @@ def test_host_bridge_scene_plan_packs_all_reads_into_one_phase_transfer( plan.close() -def test_scene_plan_pose_reads_fall_back_to_canonical_sensors_on_older_host_contracts() -> None: +def test_scene_plan_pose_reads_fall_back_to_canonical_sensors_on_older_host_contracts( + monkeypatch: pytest.MonkeyPatch, +) -> None: backend = _SceneTensorBackend() scene = _scene(backend) + _use_legacy_tensor_spec(monkeypatch) plan = scene.compile_tensor_reads( "cpu", ( @@ -463,9 +478,9 @@ def test_scene_plan_pose_reads_use_packed_pose_blocks_without_velocities( assert plan.body_fields == {"robot": ("pos_w", "quat_w")} assert plan.packet_names == ("body_pos_w", "body_quat_w") if backend.io_specs: - assert backend.io_specs[0].body_names == ("robot/hip",) + assert backend.io_specs[0].body_names == ("hip",) assert backend.io_specs[0].body_fields == ("pos_w", "quat_w") - assert plan.packet_body_axis("body_pos_w") == ("robot/hip",) + assert plan.packet_body_axis("body_pos_w") == ("hip",) plan.refresh() pose = plan.ray_pose_view("robot") diff --git a/tests/envs/mdp/test_ray_observation.py b/tests/envs/mdp/test_ray_observation.py index 14ee478fc..56fedae50 100644 --- a/tests/envs/mdp/test_ray_observation.py +++ b/tests/envs/mdp/test_ray_observation.py @@ -11,6 +11,8 @@ import numpy as np import pytest import torch +from unisim.errors import UnsupportedCapabilityError +from unisim.ray_query import RayCollisionDescription, RaySceneDescription, RayTraceOutputs from unilab.envs.mdp.ray_observation import RayPatternCfg, RayScan from unilab.managers.manager_base import ManagerTermBaseCfg @@ -23,6 +25,7 @@ class _Block: def __init__(self, values: torch.Tensor): self.distance = values + self.normal = torch.zeros(*values.shape, 3, dtype=torch.float32) @dataclass @@ -32,6 +35,9 @@ class _Capabilities: supports_device_output: bool = True stream_event_ownership: str | None = "caller" torch_devices: tuple[str, ...] = ("cpu",) + supports_min_distance: bool = True + supports_static_filter: bool = True + supports_tensor_normal: bool = True class _FakeCaster: @@ -44,6 +50,7 @@ def __init__(self, *, values: torch.Tensor | None = None): self.layout: tuple[str, ...] | None = None self.poses: list[Any] = [] self.traces = 0 + self.last_query: tuple[float, object, RayTraceOutputs] | None = None self.closed = 0 @property @@ -67,20 +74,36 @@ def bind_pose_layout(self, body_names) -> None: def update_pose_tensor(self, pose: Any) -> None: self.poses.append(pose) - def trace_tensor(self, origins, directions, max_distance, env_ids=None, outputs=None): + def trace_tensor( + self, + origins, + directions, + max_distance, + env_ids=None, + outputs=None, + *, + min_distance=0.0, + ray_filter=None, + ): self.traces += 1 self.last_inputs = (origins, directions, max_distance) + self.last_query = (min_distance, ray_filter, outputs) if origins.shape != directions.shape: raise ValueError( f"origin and direction tensor batches must have the same shape: " f"{tuple(origins.shape)} != {tuple(directions.shape)}" ) - if origins.ndim != 2 or origins.shape[-1] != 3: + if origins.ndim not in (2, 3) or origins.shape[-1] != 3: raise ValueError( - f"shared ray tensors must have shape [num_rays, 3]: {tuple(origins.shape)}" + "ray tensors must have shape [num_rays, 3] or " + f"[num_envs, num_rays, 3]: {tuple(origins.shape)}" ) - if origins.shape[0] != self.num_rays: + if origins.ndim == 2 and origins.shape[0] != self.num_rays: raise ValueError(f"expected {self.num_rays} rays, got {origins.shape[0]}") + if origins.ndim == 3 and origins.shape[:2] != (self.num_envs, self.num_rays): + raise ValueError( + f"expected {self.num_envs} x {self.num_rays} rays, got {tuple(origins.shape)}" + ) if not origins.is_contiguous() or not directions.is_contiguous(): raise ValueError("shared ray tensors must be contiguous") values = self.values @@ -92,15 +115,16 @@ def close(self) -> None: self.closed += 1 -class _RayScene: - num_bodies = 1 - num_geoms = 0 - geom_types = () - geom_sizes = np.empty((0, 3), dtype=np.float64) - geom_local_pos = np.empty((0, 3), dtype=np.float64) - geom_local_quat = np.empty((0, 4), dtype=np.float64) - geom_body_ids = np.empty((0,), dtype=np.intp) - body_names = ("base",) +def _RayScene() -> RaySceneDescription: + return RaySceneDescription( + num_bodies=2, + geom_types=("plane", "sphere"), + geom_sizes=np.array([[0.0, 0.0, 0.0], [0.5, 0.0, 0.0]], dtype=np.float64), + geom_local_pos=np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]], dtype=np.float64), + geom_local_quat=np.array([[1.0, 0.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0]], dtype=np.float64), + geom_body_ids=np.array([0, 1], dtype=np.intp), + body_names=("world", "base"), + ) class _ReadPlan: @@ -121,23 +145,38 @@ class _Backend: backend_type = "fake-ray-backend" num_envs = 2 - def __init__(self, scene: Any | None = None): + def __init__(self, scene: Any | None = None, collision: Any | None = None): self.scene = scene if scene is not None else _RayScene() + self.collision = ( + collision if collision is not None else RayCollisionDescription(scene=self.scene) + ) def get_ray_scene_description(self, body_names=None) -> Any: if tuple(body_names or ()) != self.scene.body_names: raise KeyError(tuple(body_names)) return self.scene + def get_ray_collision_description(self, body_names=None) -> Any: + if body_names is not None and tuple(body_names) != self.scene.body_names: + raise KeyError(tuple(body_names)) + return self.collision + + def get_body_ids(self, body_names) -> np.ndarray: + names = tuple(self.scene.body_names) + return np.asarray([names.index(name) for name in body_names], dtype=np.int64) + class _Env: num_envs = 2 device = torch.device("cpu") backend = _Backend() - def __init__(self, read_plan: Any, caster: _FakeCaster, scene: Any): - self.scene = SimpleNamespace(read_plan=read_plan) - self.backend = _Backend(scene) + def __init__(self, read_plan: Any, caster: _FakeCaster, scene: Any, collision: Any): + self.scene = SimpleNamespace( + read_plan=read_plan, + entities={"robot": SimpleNamespace(body_names=("base",))}, + ) + self.backend = _Backend(scene, collision) self._caster = caster @property @@ -145,52 +184,77 @@ def caster(self) -> _FakeCaster: return self._caster -def _bind(monkeypatch, values: torch.Tensor | None = None) -> tuple[_FakeCaster, RayScan]: +def _bind( + monkeypatch, + values: torch.Tensor | None = None, + pose_pos: torch.Tensor | None = None, + pose_quat: torch.Tensor | None = None, + **pattern_kwargs: Any, +) -> tuple[_FakeCaster, RayScan]: scene = _RayScene() - caster = _FakeCaster(values=values) - read_plan = _ReadPlan( - torch.randn((2, 1, 3), dtype=torch.float32), - torch.randn((2, 1, 4), dtype=torch.float32).renorm(dim=-1, p=2, maxnorm=1.0), + collision = RayCollisionDescription( + scene=scene, + geom_names=("ground", "ball"), + geom_tags=(("terrain",), ("robot",)), + body_parent_ids=np.array([-1, 0], dtype=np.intp), ) - env = _Env(read_plan, caster, scene) + caster = _FakeCaster(values=values) + if pose_pos is None: + pose_pos = torch.zeros((2, 1, 3), dtype=torch.float32) + if pose_quat is None: + pose_quat = torch.zeros((2, 1, 4), dtype=torch.float32) + pose_quat[..., 0] = 1.0 + read_plan = _ReadPlan(pose_pos, pose_quat) + env = _Env(read_plan, caster, scene, collision) pattern = RayPatternCfg( entity_name="robot", - body_names=("base",), + attachment_body_name="base", num_rays=4, pattern_start=0.0, pattern_stop=1.5, max_distance=3.0, + **pattern_kwargs, ) term = RayScan(ManagerTermBaseCfg(func=RayScan, params={"ray_pattern": pattern}), env) monkeypatch.setattr( "unilab.envs.mdp.ray_observation._create_ray_caster", - lambda caster_name, **kwargs: (setattr(caster, "scene", scene), caster)[1], + lambda caster_name, **kwargs: ( + setattr(caster, "scene", kwargs["collision"].scene), + caster, + )[1], ) term.bind_read_phase() return caster, term def test_ray_pattern_config_validates_public_names_and_values() -> None: - assert RayPatternCfg( - entity_name="robot", - body_names=("base",), - num_rays=4, - pattern_start=0, - pattern_stop=1, - max_distance=2, - ).body_names == ("base",) + assert ( + RayPatternCfg( + entity_name="robot", + attachment_body_name="base", + num_rays=4, + pattern_start=0, + pattern_stop=1, + max_distance=2, + ).attachment_body_name + == "base" + ) for kwargs in ( {"entity_name": ""}, - {"body_names": "base"}, - {"body_names": ()}, - {"body_names": ("base", "base")}, + {"attachment_body_name": None}, {"num_rays": 0}, {"max_distance": 0.0}, + {"min_distance": -0.1}, + {"min_distance": 1.0}, + {"sensor_offset": "0 0 0"}, + {"sensor_offset": (0.0, 0.0)}, + {"attachment_body_name": ""}, + {"caster_options": {"collision": object()}}, {"caster": ""}, ): base = dict( entity_name="robot", - body_names=("base",), + attachment_body_name="base", num_rays=4, pattern_start=0.0, pattern_stop=1.0, @@ -217,10 +281,9 @@ def load_module(self, fullname): result = term(term._env) assert caster.materialized - assert caster.layout == ("base",) - assert term.tensor_body_names == ("base",) + assert caster.layout == ("world", "base") assert term.observation_width == 4 - assert caster.poses[-1].body_names == ("base",) + assert caster.poses[-1].body_names == ("world", "base") assert result.shape == (2, 4) assert result.dtype == torch.float32 assert result.device == torch.device("cpu") @@ -247,27 +310,263 @@ def test_ray_scan_copies_reused_caster_output_into_term_owned_carrier(monkeypatc term.close() -def test_ray_scan_publishes_matching_shared_ray_batches(monkeypatch) -> None: +def test_ray_scan_uses_public_collision_descriptor_and_rich_query(monkeypatch) -> None: + source = torch.arange(8, dtype=torch.float32).reshape(2, 4) + caster, term = _bind( + monkeypatch, + values=source, + min_distance=0.2, + include_normal=True, + allowed_geom_names=("ball",), + allowed_tags=("robot",), + ) + result = term(term._env) + collision = term._env.backend.collision + packed = result.reshape(2, -1, 4) + + assert term.observation_width == 16 + assert result.shape == (2, term.observation_width) + assert caster.scene is collision.scene + assert caster.last_query is not None + min_distance, ray_filter, outputs = caster.last_query + assert min_distance == 0.2 + assert ray_filter is not None + assert ray_filter.geom_mask(collision).tolist() == [False, True] + assert outputs.normal + torch.testing.assert_close(packed[..., 0], source) + torch.testing.assert_close(packed[..., 1:4], torch.zeros((2, 4, 3))) + term.close() + + +def test_ray_scan_excludes_attachment_body_through_public_filter(monkeypatch) -> None: + caster, term = _bind(monkeypatch, excluded_body_names=("base",)) + term(term._env) + collision = term._env.backend.collision + + assert caster.last_query is not None + ray_filter = caster.last_query[1] + assert ray_filter is not None + assert ray_filter.geom_mask(collision).tolist() == [True, False] + term.close() + + +def test_ray_scan_fails_closed_on_unsupported_rich_capabilities(monkeypatch) -> None: + caster = _FakeCaster() + read_plan = _ReadPlan( + torch.zeros((2, 1, 3), dtype=torch.float32), + torch.zeros((2, 1, 4), dtype=torch.float32), + ) + scene = _RayScene() + collision = RayCollisionDescription( + scene=scene, + geom_names=("ground", "ball"), + geom_tags=(("terrain",), ("robot",)), + body_parent_ids=np.array([-1, 0], dtype=np.intp), + ) + env = _Env(read_plan, caster, scene, collision) + pattern = RayPatternCfg( + entity_name="robot", + attachment_body_name="base", + num_rays=2, + pattern_start=0.0, + pattern_stop=1.0, + max_distance=2.0, + min_distance=0.1, + ) + term = RayScan(ManagerTermBaseCfg(func=RayScan, params={"ray_pattern": pattern}), env) + caster.get_ray_capabilities = lambda: _Capabilities(supports_min_distance=False) + monkeypatch.setattr( + "unilab.envs.mdp.ray_observation._create_ray_caster", + lambda caster_name, **kwargs: caster, + ) + with pytest.raises(NotImplementedError, match="min_distance"): + term.bind_read_phase() + assert not caster.materialized + + +def test_ray_scan_fails_closed_on_unsupported_filter_and_normal(monkeypatch) -> None: + def bind(**pattern_kwargs): + caster = _FakeCaster() + read_plan = _ReadPlan( + torch.zeros((2, 1, 3), dtype=torch.float32), + torch.zeros((2, 1, 4), dtype=torch.float32), + ) + scene = _RayScene() + collision = RayCollisionDescription( + scene=scene, + geom_names=("ground", "ball"), + geom_tags=(("terrain",), ("robot",)), + body_parent_ids=np.array([-1, 0], dtype=np.intp), + ) + env = _Env(read_plan, caster, scene, collision) + values = { + "entity_name": "robot", + "attachment_body_name": "base", + "num_rays": 2, + "pattern_start": 0.0, + "pattern_stop": 1.0, + "max_distance": 2.0, + } + values.update(pattern_kwargs) + term = RayScan( + ManagerTermBaseCfg(func=RayScan, params={"ray_pattern": RayPatternCfg(**values)}), + env, + ) + monkeypatch.setattr( + "unilab.envs.mdp.ray_observation._create_ray_caster", + lambda caster_name, **kwargs: caster, + ) + return caster, term + + caster, term = bind(excluded_body_names=("base",)) + caster.get_ray_capabilities = lambda: _Capabilities(supports_static_filter=False) + with pytest.raises(NotImplementedError, match="static ray filters"): + term.bind_read_phase() + + caster, term = bind(include_normal=True) + caster.get_ray_capabilities = lambda: _Capabilities(supports_tensor_normal=False) + with pytest.raises(NotImplementedError, match="tensor normals"): + term.bind_read_phase() + + +def test_ray_scan_maps_ancestor_exclusion_from_public_names(monkeypatch) -> None: + caster, term = _bind(monkeypatch, excluded_ancestor_body_names=("world",)) + term(term._env) + collision = term._env.backend.collision + + assert caster.last_query is not None + ray_filter = caster.last_query[1] + assert ray_filter is not None + assert ray_filter.geom_mask(collision).tolist() == [False, False] + term.close() + + +def test_ray_scan_fails_closed_on_missing_ancestry(monkeypatch) -> None: + caster = _FakeCaster() + read_plan = _ReadPlan( + torch.zeros((2, 1, 3), dtype=torch.float32), + torch.zeros((2, 1, 4), dtype=torch.float32), + ) + scene = _RayScene() + collision = RayCollisionDescription( + scene=scene, + geom_names=("ground", "ball"), + geom_tags=(("terrain",), ("robot",)), + ) + env = _Env(read_plan, caster, scene, collision) + pattern = RayPatternCfg( + entity_name="robot", + attachment_body_name="base", + num_rays=2, + pattern_start=0.0, + pattern_stop=1.0, + max_distance=2.0, + excluded_ancestor_body_names=("world",), + ) + term = RayScan(ManagerTermBaseCfg(func=RayScan, params={"ray_pattern": pattern}), env) + with pytest.raises(UnsupportedCapabilityError, match="ancestor"): + term.bind_read_phase() + + +def test_ray_scan_fails_closed_on_uncovered_non_root_body(monkeypatch) -> None: + scene = RaySceneDescription( + num_bodies=3, + geom_types=("plane", "sphere", "box"), + geom_sizes=np.array([[0.0, 0.0, 0.0], [0.5, 0.0, 0.0], [0.2, 0.2, 0.2]], dtype=np.float64), + geom_local_pos=np.zeros((3, 3), dtype=np.float64), + geom_local_quat=np.tile(np.array([1.0, 0.0, 0.0, 0.0]), (3, 1)), + geom_body_ids=np.array([0, 1, 2], dtype=np.intp), + body_names=("world", "base", "static_object"), + ) + collision = RayCollisionDescription( + scene=scene, + geom_names=("ground", "ball", "object"), + geom_tags=(("terrain",), ("robot",), ("terrain",)), + body_parent_ids=np.array([-1, 0, 0], dtype=np.intp), + ) + caster = _FakeCaster() + read_plan = _ReadPlan( + torch.zeros((2, 1, 3), dtype=torch.float32), + torch.zeros((2, 1, 4), dtype=torch.float32), + ) + env = _Env(read_plan, caster, scene, collision) + pattern = RayPatternCfg( + entity_name="robot", + attachment_body_name="base", + num_rays=2, + pattern_start=0.0, + pattern_stop=1.0, + max_distance=2.0, + ) + term = RayScan(ManagerTermBaseCfg(func=RayScan, params={"ray_pattern": pattern}), env) + with pytest.raises(NotImplementedError, match="uncovered bodies"): + term.bind_read_phase() + + +def test_ray_scan_publishes_matching_attachment_ray_batches(monkeypatch) -> None: caster, term = _bind(monkeypatch) term(term._env) origins, directions, max_distance = caster.last_inputs - assert origins.shape == directions.shape == (4, 3) + assert origins.shape == directions.shape == (2, 4, 3) assert origins.is_contiguous() and directions.is_contiguous() assert origins.dtype == directions.dtype == torch.float32 assert origins.device == directions.device == torch.device("cpu") torch.testing.assert_close( origins, torch.tensor( - [[0.0, 0.0, term.ray_pattern.pattern_center_height]] * 4, + [[[0.0, 0.0, 0.0]] * 4] * 2, dtype=torch.float32, ), ) + angles = torch.linspace( + term.ray_pattern.pattern_start, + term.ray_pattern.pattern_stop, + term.ray_pattern.num_rays, + dtype=torch.float32, + ) + expected_directions = torch.stack( + (torch.cos(angles), torch.sin(angles), torch.zeros_like(angles)), dim=-1 + ).repeat(2, 1, 1) + torch.testing.assert_close(directions, expected_directions) assert max_distance == term.ray_pattern.max_distance term.close() +def test_ray_pattern_follows_body_attachment_pose(monkeypatch) -> None: + pose_pos = torch.tensor([[[0.1, 0.2, 0.3]], [[-0.4, 0.5, -0.2]]], dtype=torch.float32) + pose_quat = torch.zeros((2, 1, 4), dtype=torch.float32) + pose_quat[..., 0] = 0.5**0.5 + pose_quat[..., 3] = 0.5**0.5 + caster, term = _bind( + monkeypatch, + pose_pos=pose_pos, + pose_quat=pose_quat, + sensor_offset=(0.2, 0.0, 0.1), + ) + + term(term._env) + origins, directions, _ = caster.last_inputs + pose = caster.poses[-1] + assert pose.body_names == ("world", "base") + assert torch.count_nonzero(pose.pos_w[:, 0, :]).item() == 0 + torch.testing.assert_close(pose.pos_w[:, 1, :], pose_pos[:, 0, :]) + rotation = torch.tensor( + [[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]], + dtype=torch.float32, + ) + expected_origin = pose_pos + torch.tensor([0.0, 0.2, 0.1], dtype=torch.float32)[None, None, :] + angles = torch.linspace(0.0, 1.5, 4, dtype=torch.float32) + local_directions = torch.stack( + (torch.cos(angles), torch.sin(angles), torch.zeros_like(angles)), dim=-1 + ) + expected_directions = local_directions[None, :, :] @ rotation.T + torch.testing.assert_close(origins, expected_origin.repeat(1, 4, 1)) + torch.testing.assert_close(directions, expected_directions.repeat(2, 1, 1)) + term.close() + + def test_ray_scan_fails_closed_before_binding_and_on_unsupported_capability(monkeypatch) -> None: scene = _RayScene() caster = _FakeCaster() @@ -275,10 +574,16 @@ def test_ray_scan_fails_closed_before_binding_and_on_unsupported_capability(monk torch.zeros((2, 1, 3), dtype=torch.float32), torch.zeros((2, 1, 4), dtype=torch.float32), ) - env = _Env(read_plan, caster, scene) + collision = RayCollisionDescription( + scene=scene, + geom_names=("ground", "ball"), + geom_tags=(("terrain",), ("robot",)), + body_parent_ids=np.array([-1, 0], dtype=np.intp), + ) + env = _Env(read_plan, caster, scene, collision) pattern = RayPatternCfg( entity_name="robot", - body_names=("base",), + attachment_body_name="base", num_rays=2, pattern_start=0.0, pattern_stop=1.0, @@ -303,9 +608,19 @@ def test_ray_scan_uses_only_public_backend_and_factory_boundaries(monkeypatch) - caster, term = _bind(monkeypatch) env = term._env - scene = env.backend.get_ray_scene_description(("base",)) - assert scene is env.backend.scene + collision = env.backend.get_ray_collision_description() + assert collision.scene is env.backend.scene + assert collision is env.backend.collision assert term.ray_pattern.caster == "uni_ray" + source = (_REPO_ROOT / "src/unilab/envs/mdp/ray_observation.py").read_text(encoding="utf-8") + for banned in ( + "import uni_ray", + "MjModel", + "mjwarp.Model", + "MjBatchCollision", + "build_collision_description", + ): + assert banned not in source, banned assert not any( path.name.startswith("_") and path.is_file() for path in Path(__file__).parent.rglob("*.py") ) diff --git a/tests/envs/test_manager_based_rl_env.py b/tests/envs/test_manager_based_rl_env.py index 25382bb2c..8f28fe0ba 100644 --- a/tests/envs/test_manager_based_rl_env.py +++ b/tests/envs/test_manager_based_rl_env.py @@ -1580,6 +1580,93 @@ def test_mujoco_host_bridge_packed_reset_lifecycle( backend.close() +def test_mujoco_tensor_ray_observation_uses_public_rich_query_contract( + tmp_path: Path, +) -> None: + """Exercise the real MuJoCo/UniSim/uni_ray perception path end to end.""" + + pytest.importorskip("uni_ray", reason="local uni_ray is required") + pytest.importorskip("warp", reason="warp-lang is required by uni_ray") + + def build(pattern: mdp.RayPatternCfg): + model = tmp_path / "ray-rich-query.xml" + model.write_text( + "" + "" + "" + "" + ) + scene = SceneCfg( + model_file=str(model), + entities={ + "robot": EntityCfg( + root_body_name="base", + body_names=("base",), + joint_names=("joint",), + actuator_names=("motor",), + ) + }, + ) + cfg = ManagerBasedRlEnvCfg( + scene=scene, + sim_dt=0.01, + ctrl_dt=0.02, + max_episode_seconds=0.04, + observations={ + "actor": ObservationGroupCfg( + terms={ + "rays": ObservationTermCfg( + func=mdp.RayScan, params={"ray_pattern": pattern} + ) + } + ) + }, + actions={ + "effort": mdp.JointEffortActionCfg(entity_name="robot", actuator_names=(".*",)) + }, + events={"default": EventTermCfg(func=mdp.reset_scene_to_default, mode="reset")}, + rewards={"alive": RewardTermCfg(func=mdp.is_alive, weight=1.0)}, + terminations={"time_out": TerminationTermCfg(func=mdp.time_out, time_out=True)}, + policy_observation_group="actor", + ) + backend = create_backend( + "mujoco", scene, 2, cfg.sim_dt, base_name="base", add_body_sensors=True + ) + return ManagerBasedRlEnv(cfg, backend, 2), backend + + def ray_pattern(**kwargs): + values = { + "entity_name": "robot", + "attachment_body_name": "base", + "num_rays": 2, + "pattern_start": 0.0, + "pattern_stop": 0.0, + "pattern_pitch": -0.5 * 3.141592653589793, + "max_distance": 3.0, + "sensor_offset": (0.0, 0.0, 0.2), + "include_normal": True, + } + values.update(kwargs) + return mdp.RayPatternCfg(**values) + + near_env, near_backend = build(ray_pattern()) + valid_env, valid_backend = build(ray_pattern(min_distance=0.15, excluded_body_names=("base",))) + try: + near = near_env.init_state().obs["obs"].reshape(2, -1, 4) + valid = valid_env.init_state().obs["obs"].reshape(2, -1, 4) + torch.testing.assert_close(near[..., 0], torch.full((2, 2), 0.1), atol=2e-5, rtol=0.0) + expected_normal = torch.tensor([0.0, 0.0, 1.0]).repeat(2, 2, 1) + torch.testing.assert_close(near[..., 1:4], expected_normal, atol=2e-5, rtol=0.0) + torch.testing.assert_close(valid[..., 0], torch.full((2, 2), 1.2), atol=2e-5, rtol=0.0) + torch.testing.assert_close(valid[..., 1:4], expected_normal, atol=2e-5, rtol=0.0) + finally: + near_env.close() + valid_env.close() + near_backend.close() + valid_backend.close() + + def _make_state_env( *, commands: dict[str, CommandTermCfg | None] | None = None, diff --git a/tests/scripts/test_rocm_dependency_profile.py b/tests/scripts/test_rocm_dependency_profile.py index ed954c2d8..6e8197a47 100644 --- a/tests/scripts/test_rocm_dependency_profile.py +++ b/tests/scripts/test_rocm_dependency_profile.py @@ -65,7 +65,7 @@ def test_ray_local_optional_dependencies_stay_branch_local() -> None: for path in (ROOT / "pyproject.toml", ROOT / "pyproject.rocm.toml"): project = _project(path) extras = project["project"]["optional-dependencies"] - assert extras["ray-local"] == ["uni-ray"] + assert extras["ray-local"] == ["uni-ray[warp]"] assert "uni-ray" not in _dependencies(path) sources = project["tool"]["uv"]["sources"] diff --git a/uv.lock b/uv.lock index 3402bb6c2..8f4563a6e 100644 --- a/uv.lock +++ b/uv.lock @@ -5153,6 +5153,44 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/dc/9b/47798a6c91d8bdb567fe2698fe81e0c6b7cb7ef4d13da4114b41d239f65d/typing_inspection-0.4.2-py3-none-any.whl", hash = "sha256:4ed1cacbdc298c220f1bd249ed5287caa16f34d44ef4e9c3d0cbad5b521545e7", size = 14611, upload-time = "2025-10-01T02:14:40.154Z" }, ] +[[package]] +name = "uni-ray" +version = "0.1.0" +source = { editable = "../uni_ray" } +dependencies = [ + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "unisim-core" }, +] + +[package.optional-dependencies] +warp = [ + { name = "warp-lang" }, +] + +[package.metadata] +requires-dist = [ + { name = "motrixsim-core", marker = "python_full_version == '3.10.*' and extra == 'motrix'", specifier = "==0.10.1" }, + { name = "mujoco", marker = "extra == 'mujoco'", specifier = ">=3.2.0" }, + { name = "numpy" }, + { name = "unisim-core", editable = "../unisim" }, + { name = "viser", marker = "extra == 'viz'" }, + { name = "warp-lang", marker = "extra == 'warp'", specifier = ">=1.11.0" }, +] +provides-extras = ["warp", "mujoco", "motrix", "viz"] + +[package.metadata.requires-dev] +dev = [ + { name = "mujoco", specifier = ">=3.2.0" }, + { name = "mujoco-warp", specifier = "~=3.11.0" }, + { name = "mypy" }, + { name = "pytest" }, + { name = "ruff" }, + { name = "torch" }, + { name = "viser" }, + { name = "warp-lang", specifier = ">=1.11.0" }, +] + [[package]] name = "unilab" version = "1.5.1" @@ -5217,6 +5255,10 @@ newton = [ { name = "pyglet" }, { name = "warp-lang" }, ] +ray-local = [ + { name = "uni-ray", extra = ["warp"] }, + { name = "unisim-core" }, +] superdex = [ { name = "unisim-core", extra = ["superdex"], marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, ] @@ -5268,15 +5310,19 @@ requires-dist = [ { name = "tqdm" }, { name = "trimesh", specifier = ">=3.21.7" }, { name = "typing-extensions" }, - { name = "unisim-core", specifier = "==2.1.1" }, - { name = "unisim-core", extras = ["motrix"], marker = "extra == 'motrix'", specifier = "==2.1.1" }, - { name = "unisim-core", extras = ["superdex"], marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'superdex'", specifier = "==2.1.1" }, + { name = "uni-ray", extras = ["warp"], marker = "extra == 'ray-local'", editable = "../uni_ray" }, + { name = "unisim-core", marker = "extra == 'ray-local'", editable = "../unisim" }, + { name = "unisim-core", marker = "extra != 'ray-local'", specifier = ">=2.1.1" }, + { name = "unisim-core", extras = ["motrix"], marker = "extra == 'motrix' and extra == 'ray-local'", editable = "../unisim" }, + { name = "unisim-core", extras = ["motrix"], marker = "extra == 'motrix' and extra != 'ray-local'", specifier = ">=2.1.1" }, + { name = "unisim-core", extras = ["superdex"], marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux' and extra != 'ray-local' and extra == 'superdex'", specifier = ">=2.1.1" }, + { name = "unisim-core", extras = ["superdex"], marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'ray-local' and extra == 'superdex'", editable = "../unisim" }, { name = "viser", specifier = ">=1.0.26" }, { name = "wandb", specifier = ">=0.22.3" }, { name = "warp-lang", marker = "extra == 'mjwarp'", specifier = "==1.16.0" }, { name = "warp-lang", marker = "extra == 'newton'", specifier = "==1.16.0" }, ] -provides-extras = ["drake", "mujoco", "mjwarp", "newton", "motrix", "genesis", "superdex"] +provides-extras = ["drake", "mujoco", "mjwarp", "newton", "motrix", "genesis", "superdex", "ray-local"] [package.metadata.requires-dev] dev = [ @@ -5289,16 +5335,12 @@ dev = [ [[package]] name = "unisim-core" -version = "2.1.1" -source = { registry = "https://pypi.org/simple" } +version = "2.2.0" +source = { editable = "../unisim" } dependencies = [ { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/ae/dc/5c8b56a04996cc3f47767ac4c3c665c0eae131a188f1554525ff7514a9ed/unisim_core-2.1.1.tar.gz", hash = "sha256:f3c9bc2497b3ef85ec984cf38f8a8ce5cfd24b900db2734685575c5bb8cef2ed", size = 699580, upload-time = "2026-10-22T10:26:44.017Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/31/99/b532c452a8544eb508ff3b6009146aa24f44e09c5c43e63aaf46153b8490/unisim_core-2.1.1-py3-none-any.whl", hash = "sha256:fee2dcecb8ceb60ee2532a9e7d15f66daf1c1213d990fe4a607aef92130f6bf2", size = 746595, upload-time = "2026-10-22T10:26:42.255Z" }, -] [package.optional-dependencies] motrix = [ @@ -5310,6 +5352,37 @@ superdex = [ { name = "superdex-robotics-uni" }, ] +[package.metadata] +requires-dist = [ + { name = "drake-uni", marker = "extra == 'drake'", specifier = "==0.1.0" }, + { name = "genesis-world", marker = "extra == 'genesis'", specifier = "==1.3.3" }, + { name = "imgui-bundle", marker = "extra == 'newton'", specifier = ">=1.92.0" }, + { name = "mjbatch-uni", marker = "extra == 'mujoco'", specifier = "~=0.2.5" }, + { name = "motrixsim-core", marker = "extra == 'motrix'", specifier = "==0.10.1" }, + { name = "mujoco", marker = "python_full_version >= '3.12' and python_full_version < '3.14' and extra == 'superdex'", specifier = "~=3.11.0" }, + { name = "mujoco", marker = "extra == 'mujoco'", specifier = "~=3.11.0" }, + { name = "mujoco", marker = "extra == 'newton'", specifier = "==3.11.0" }, + { name = "mujoco", marker = "extra == 'scene-compiler'", specifier = "~=3.11.0" }, + { name = "mujoco-warp", marker = "extra == 'mjwarp'", specifier = "~=3.11.0" }, + { name = "mujoco-warp", marker = "extra == 'newton'", specifier = "==3.11.0" }, + { name = "newton", marker = "extra == 'newton'", specifier = "==1.5.1" }, + { name = "numpy" }, + { name = "pyglet", marker = "extra == 'newton'", specifier = ">=2.1.6,<3" }, + { name = "superdex-physics-uni", marker = "python_full_version >= '3.12' and python_full_version < '3.14' and extra == 'superdex'", specifier = "==1.3.0" }, + { name = "superdex-robotics-uni", marker = "python_full_version >= '3.12' and python_full_version < '3.14' and extra == 'superdex'", specifier = "==1.3.0" }, + { name = "warp-lang", marker = "extra == 'mjwarp'", specifier = "==1.16.0" }, + { name = "warp-lang", marker = "extra == 'newton'", specifier = "==1.16.0" }, +] +provides-extras = ["scene-compiler", "mujoco", "motrix", "drake", "mjwarp", "genesis", "superdex", "newton", "isaacgym", "isaacsim"] + +[package.metadata.requires-dev] +dev = [ + { name = "mypy" }, + { name = "pyright", specifier = "==1.1.408" }, + { name = "pytest" }, + { name = "ruff" }, +] + [[package]] name = "unrealcv" version = "1.3.2" diff --git a/uv.rocm.lock b/uv.rocm.lock index fcde2bf27..9a69270d6 100644 --- a/uv.rocm.lock +++ b/uv.rocm.lock @@ -3707,6 +3707,44 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/dc/9b/47798a6c91d8bdb567fe2698fe81e0c6b7cb7ef4d13da4114b41d239f65d/typing_inspection-0.4.2-py3-none-any.whl", hash = "sha256:4ed1cacbdc298c220f1bd249ed5287caa16f34d44ef4e9c3d0cbad5b521545e7", size = 14611, upload-time = "2025-10-01T02:14:40.154Z" }, ] +[[package]] +name = "uni-ray" +version = "0.1.0" +source = { editable = "../uni_ray" } +dependencies = [ + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "unisim-core" }, +] + +[package.optional-dependencies] +warp = [ + { name = "warp-lang" }, +] + +[package.metadata] +requires-dist = [ + { name = "motrixsim-core", marker = "python_full_version == '3.10.*' and extra == 'motrix'", specifier = "==0.10.1" }, + { name = "mujoco", marker = "extra == 'mujoco'", specifier = ">=3.2.0" }, + { name = "numpy" }, + { name = "unisim-core", editable = "../unisim" }, + { name = "viser", marker = "extra == 'viz'" }, + { name = "warp-lang", marker = "extra == 'warp'", specifier = ">=1.11.0" }, +] +provides-extras = ["warp", "mujoco", "motrix", "viz"] + +[package.metadata.requires-dev] +dev = [ + { name = "mujoco", specifier = ">=3.2.0" }, + { name = "mujoco-warp", specifier = "~=3.11.0" }, + { name = "mypy" }, + { name = "pytest" }, + { name = "ruff" }, + { name = "torch" }, + { name = "viser" }, + { name = "warp-lang", specifier = ">=1.11.0" }, +] + [[package]] name = "unilab" version = "1.5.1" @@ -3756,6 +3794,10 @@ mujoco = [ { name = "mjbatch-uni" }, { name = "mujoco" }, ] +ray-local = [ + { name = "uni-ray", extra = ["warp"] }, + { name = "unisim-core" }, +] [package.dev-dependencies] dev = [ @@ -3797,12 +3839,15 @@ requires-dist = [ { name = "trimesh", specifier = ">=3.21.7" }, { name = "triton-rocm", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'", specifier = "==3.8.0", index = "https://download.pytorch.org/whl/rocm7.2" }, { name = "typing-extensions" }, - { name = "unisim-core", specifier = "==2.1.1" }, - { name = "unisim-core", extras = ["motrix"], marker = "extra == 'motrix'", specifier = "==2.1.1" }, + { name = "uni-ray", extras = ["warp"], marker = "extra == 'ray-local'", editable = "../uni_ray" }, + { name = "unisim-core", marker = "extra == 'ray-local'", editable = "../unisim" }, + { name = "unisim-core", marker = "extra != 'ray-local'", specifier = ">=2.1.1" }, + { name = "unisim-core", extras = ["motrix"], marker = "extra == 'motrix' and extra == 'ray-local'", editable = "../unisim" }, + { name = "unisim-core", extras = ["motrix"], marker = "extra == 'motrix' and extra != 'ray-local'", specifier = ">=2.1.1" }, { name = "viser", specifier = ">=1.0.26" }, { name = "wandb", specifier = ">=0.22.3" }, ] -provides-extras = ["drake", "mujoco", "motrix"] +provides-extras = ["drake", "mujoco", "motrix", "ray-local"] [package.metadata.requires-dev] dev = [ @@ -3815,22 +3860,49 @@ dev = [ [[package]] name = "unisim-core" -version = "2.1.1" -source = { registry = "https://pypi.org/simple" } +version = "2.2.0" +source = { editable = "../unisim" } dependencies = [ { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/ae/dc/5c8b56a04996cc3f47767ac4c3c665c0eae131a188f1554525ff7514a9ed/unisim_core-2.1.1.tar.gz", hash = "sha256:f3c9bc2497b3ef85ec984cf38f8a8ce5cfd24b900db2734685575c5bb8cef2ed", size = 699580, upload-time = "2026-10-22T10:26:44.017Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/31/99/b532c452a8544eb508ff3b6009146aa24f44e09c5c43e63aaf46153b8490/unisim_core-2.1.1-py3-none-any.whl", hash = "sha256:fee2dcecb8ceb60ee2532a9e7d15f66daf1c1213d990fe4a607aef92130f6bf2", size = 746595, upload-time = "2026-10-22T10:26:42.255Z" }, -] [package.optional-dependencies] motrix = [ { name = "motrixsim-core" }, ] +[package.metadata] +requires-dist = [ + { name = "drake-uni", marker = "extra == 'drake'", specifier = "==0.1.0" }, + { name = "genesis-world", marker = "extra == 'genesis'", specifier = "==1.3.3" }, + { name = "imgui-bundle", marker = "extra == 'newton'", specifier = ">=1.92.0" }, + { name = "mjbatch-uni", marker = "extra == 'mujoco'", specifier = "~=0.2.5" }, + { name = "motrixsim-core", marker = "extra == 'motrix'", specifier = "==0.10.1" }, + { name = "mujoco", marker = "python_full_version >= '3.12' and python_full_version < '3.14' and extra == 'superdex'", specifier = "~=3.11.0" }, + { name = "mujoco", marker = "extra == 'mujoco'", specifier = "~=3.11.0" }, + { name = "mujoco", marker = "extra == 'newton'", specifier = "==3.11.0" }, + { name = "mujoco", marker = "extra == 'scene-compiler'", specifier = "~=3.11.0" }, + { name = "mujoco-warp", marker = "extra == 'mjwarp'", specifier = "~=3.11.0" }, + { name = "mujoco-warp", marker = "extra == 'newton'", specifier = "==3.11.0" }, + { name = "newton", marker = "extra == 'newton'", specifier = "==1.5.1" }, + { name = "numpy" }, + { name = "pyglet", marker = "extra == 'newton'", specifier = ">=2.1.6,<3" }, + { name = "superdex-physics-uni", marker = "python_full_version >= '3.12' and python_full_version < '3.14' and extra == 'superdex'", specifier = "==1.3.0" }, + { name = "superdex-robotics-uni", marker = "python_full_version >= '3.12' and python_full_version < '3.14' and extra == 'superdex'", specifier = "==1.3.0" }, + { name = "warp-lang", marker = "extra == 'mjwarp'", specifier = "==1.16.0" }, + { name = "warp-lang", marker = "extra == 'newton'", specifier = "==1.16.0" }, +] +provides-extras = ["scene-compiler", "mujoco", "motrix", "drake", "mjwarp", "genesis", "superdex", "newton", "isaacgym", "isaacsim"] + +[package.metadata.requires-dev] +dev = [ + { name = "mypy" }, + { name = "pyright", specifier = "==1.1.408" }, + { name = "pytest" }, + { name = "ruff" }, +] + [[package]] name = "urllib3" version = "2.6.3" @@ -3932,6 +4004,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c4/e8/76836b75d401ff5912aaf513176e64557ceaec4c4946bfd38a698ff84d48/wandb-0.25.1-py3-none-win_arm64.whl", hash = "sha256:cc7c34b70cf4b7be4d395541e82e325fd9d2be978d62c9ec01f1a7141523b6bb", size = 22080774, upload-time = "2026-03-10T23:51:40.196Z" }, ] +[[package]] +name = "warp-lang" +version = "1.18.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/59/ca/0f88f27e9367b55df2c87fd6dbb96f15871ab7fcca14dd9ec92ec4688a30/warp_lang-1.18.0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:3d0a12f411555c08015f84060187a0466800d2c837eb8800977b9604be438c21", size = 27819487, upload-time = "2026-10-05T00:22:32.346Z" }, + { url = "https://files.pythonhosted.org/packages/13/69/5fc4388c6af4c78cd247776bb6928cabe83d893896444dd2127ccbe8b2ca/warp_lang-1.18.0-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:090962bfae7dd22ade0d8e0258f41635e0e58980f15bcc9e7a7d312819dfe440", size = 148811186, upload-time = "2026-10-05T00:23:01.197Z" }, + { url = "https://files.pythonhosted.org/packages/be/41/6e255904a10cd1b3b1a61cc9ac21551812ffc55b1c68324de4c0c24a6660/warp_lang-1.18.0-py3-none-manylinux_2_34_aarch64.whl", hash = "sha256:d76ab8c7d4f4897405fa0a744d38812fc44cfc8162c40e658dc9ca7790b6a645", size = 147348850, upload-time = "2026-10-05T00:23:31.015Z" }, + { url = "https://files.pythonhosted.org/packages/3e/75/888dfe20326fea1a3b46b68a239ccd5db717bb039d511ab6fca7c77e6fea/warp_lang-1.18.0-py3-none-win_amd64.whl", hash = "sha256:1769e55746739f2e244a2b412123003c462ec3d977dfb5f420f924a8553d30aa", size = 134607061, upload-time = "2026-10-05T00:23:56.792Z" }, +] + [[package]] name = "wcwidth" version = "0.6.0"