diff --git a/docs/sphinx/source/changelog.md b/docs/sphinx/source/changelog.md index 11da67e8f..72033eb08 100644 --- a/docs/sphinx/source/changelog.md +++ b/docs/sphinx/source/changelog.md @@ -15,6 +15,48 @@ PyPI 版本变更与未发布变更；发布日期采用 PyPI 上传日期。完 ## 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 负责。 + ### Changed / 变更 - Genesis runtime domain randomization now follows the unified device-resident diff --git a/pyproject.rocm.toml b/pyproject.rocm.toml index 6ff905e47..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,7 +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. +# 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 @@ -91,6 +97,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..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,13 +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. +# 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 @@ -129,6 +136,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 688063ce3..173e80fe7 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 @@ -206,11 +207,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], ...] = ( @@ -219,23 +235,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: @@ -244,8 +288,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 @@ -256,8 +303,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: @@ -278,6 +336,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 @@ -285,20 +348,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 @@ -323,6 +390,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 @@ -337,11 +414,53 @@ 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() 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 = 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() self._publish_packet(packet) @@ -355,7 +474,17 @@ 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 = 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() self._publish_packet(packet) @@ -371,6 +500,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() @@ -501,14 +646,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( @@ -520,8 +660,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 @@ -554,8 +717,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], @@ -569,15 +733,123 @@ 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)) + 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) + 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: @@ -633,18 +905,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 @@ -664,9 +950,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 " @@ -707,17 +1007,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 @@ -4253,6 +4574,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: @@ -4265,6 +4587,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: @@ -4275,6 +4598,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(): @@ -4294,17 +4619,68 @@ 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 + ) + 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() + ) + 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 + # 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": ( + sensor_packet_names if host_sensor_packet_names is not None else packet_names + ), + "device": resolved_device, + } + if host_pose_only_blocks and any(expanded_bodies.values()): + tensor_io_options.update( + { + "body_names": native_body_names, + "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: @@ -4319,6 +4695,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, ) @@ -4372,9 +4760,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] = [] @@ -4384,13 +4773,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 @@ -4495,5 +4903,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 5df33fd5a..5253fe4d0 100644 --- a/src/unilab/envs/manager_based_rl_env.py +++ b/src/unilab/envs/manager_based_rl_env.py @@ -430,12 +430,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()) @@ -510,6 +512,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) @@ -535,6 +541,41 @@ 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] = [] + 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) + 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__}" + ) + 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]: specs: list[SceneTensorReadSpec] = [] narrow = ( @@ -1728,6 +1769,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..409823698 --- /dev/null +++ b/src/unilab/envs/mdp/ray_observation.py @@ -0,0 +1,479 @@ +"""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 unisim.ray_query import RayPoseTensorView, RayTraceFilter, RayTraceOutputs + +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. ``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 + 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" + 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.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), + ("pattern_pitch", self.pattern_pitch), + ("max_distance", self.max_distance), + ("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") + 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 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") + 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 a distance/normal 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.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._caster: Any = None + self._bound = False + self._closed = False + self._distance: 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.""" + 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: + 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, + ) + 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, + self.ray_pattern.num_rays, + dtype=torch.float32, + device=self._device, + ) + 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._local_origins = torch.zeros( + (self.ray_pattern.num_rays, 3), + dtype=torch.float32, + device=self._device, + ) + 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, + 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 = self._refresh_full_pose(read_plan) + try: + self._refresh_attachment_rays(pose) + self._caster.update_pose_tensor(pose) + result = self._caster.trace_tensor( + 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 + 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.""" + 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 _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 + + 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..79b0ddce3 100644 --- a/tests/base/test_entity_tensor_read_plan.py +++ b/tests/base/test_entity_tensor_read_plan.py @@ -2,7 +2,8 @@ from __future__ import annotations -from types import MethodType +from dataclasses import dataclass, field +from types import MethodType, SimpleNamespace from typing import cast import numpy as np @@ -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: @@ -208,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 @@ -226,6 +244,8 @@ 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 return packet @@ -242,6 +262,30 @@ 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 _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",)), @@ -249,8 +293,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 +323,32 @@ 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( + monkeypatch: pytest.MonkeyPatch, +) -> None: + backend = _SceneTensorBackend() + scene = _scene(backend) + _use_legacy_tensor_spec(monkeypatch) + 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 +404,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 +423,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 +460,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 == ("hip",) + assert backend.io_specs[0].body_fields == ("pos_w", "quat_w") + assert plan.packet_body_axis("body_pos_w") == ("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..56fedae50 --- /dev/null +++ b/tests/envs/mdp/test_ray_observation.py @@ -0,0 +1,627 @@ +"""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 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 + +_REPO_ROOT = Path(__file__).resolve().parents[3] + + +class _Block: + """Minimal RayTraceTensorResult-shaped caster output.""" + + def __init__(self, values: torch.Tensor): + self.distance = values + self.normal = torch.zeros(*values.shape, 3, dtype=torch.float32) + + +@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",) + supports_min_distance: bool = True + supports_static_filter: bool = True + supports_tensor_normal: bool = True + + +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.last_query: tuple[float, object, RayTraceOutputs] | None = None + 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, + *, + 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 not in (2, 3) or origins.shape[-1] != 3: + raise ValueError( + "ray tensors must have shape [num_rays, 3] or " + f"[num_envs, num_rays, 3]: {tuple(origins.shape)}" + ) + 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 + if values is None: + values = torch.arange(8, dtype=torch.float32).reshape(2, 4) + return _Block(values) + + def close(self) -> None: + self.closed += 1 + + +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: + 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, 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, collision: Any): + self.scene = SimpleNamespace( + read_plan=read_plan, + entities={"robot": SimpleNamespace(body_names=("base",))}, + ) + self.backend = _Backend(scene, collision) + self._caster = caster + + @property + def caster(self) -> _FakeCaster: + return self._caster + + +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() + collision = RayCollisionDescription( + scene=scene, + geom_names=("ground", "ball"), + geom_tags=(("terrain",), ("robot",)), + body_parent_ids=np.array([-1, 0], dtype=np.intp), + ) + 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", + 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", 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", + 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": ""}, + {"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", + attachment_body_name="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 == ("world", "base") + assert term.observation_width == 4 + assert caster.poses[-1].body_names == ("world", "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_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 == (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, 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() + read_plan = _ReadPlan( + torch.zeros((2, 1, 3), dtype=torch.float32), + torch.zeros((2, 1, 4), dtype=torch.float32), + ) + 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=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 + + 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") + ) + term.close() 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/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..6e8197a47 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[warp]"] + 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'"} + ] 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"