# ----------------------------------------------------------------------------
# Copyright (c) 2021-2026 DexForce Technology Co., Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ----------------------------------------------------------------------------
"""Observed robot state, symbolic task state, and scene snapshots."""
from __future__ import annotations
import math
from collections.abc import Iterator, Mapping
from dataclasses import dataclass, field
from types import MappingProxyType
from typing import TYPE_CHECKING
import torch
if TYPE_CHECKING:
from .core import ObjectSemantics
def _same_physical_object(
first: ObjectSemantics,
second: ObjectSemantics,
) -> bool:
"""Return whether two semantic records identify one physical object."""
from .core import _same_object_identity
return _same_object_identity(first, second)
def _resolve_runtime_device(device: torch.device | str) -> torch.device:
"""Resolve an indexless CUDA device to the active concrete GPU index."""
resolved = torch.device(device)
if resolved.type == "cuda" and resolved.index is None:
return torch.device(f"cuda:{torch.cuda.current_device()}")
return resolved
def _validate_pose(value: torch.Tensor, name: str) -> int | None:
"""Validate a homogeneous transform and return its explicit batch size."""
if not isinstance(value, torch.Tensor):
raise TypeError(f"{name} must be a torch.Tensor.")
if value.shape == (4, 4):
return None
if value.dim() != 3 or value.shape[-2:] != (4, 4) or value.shape[0] == 0:
raise ValueError(
f"{name} must have shape (4, 4) or (num_envs, 4, 4), "
f"got {tuple(value.shape)}."
)
return int(value.shape[0])
def _normalize_mask(
value: torch.Tensor | None,
*,
batch_size: int,
device: torch.device,
name: str,
) -> torch.Tensor:
"""Return an owned boolean mask with shape ``(batch_size,)``."""
if value is None:
return torch.ones(batch_size, dtype=torch.bool, device=device)
if not isinstance(value, torch.Tensor):
raise TypeError(f"{name} must be a torch.Tensor or None.")
if value.dtype != torch.bool:
raise TypeError(f"{name} must have dtype torch.bool, got {value.dtype}.")
if value.shape != (batch_size,):
raise ValueError(
f"{name} must have shape ({batch_size},), got {tuple(value.shape)}."
)
return value.to(device=device).clone()
def _broadcast_pose(
value: torch.Tensor,
*,
batch_size: int,
device: torch.device,
name: str,
) -> torch.Tensor:
"""Resolve an optionally batched pose to the task-state batch."""
pose_batch_size = _validate_pose(value, name)
if value.device != device:
raise ValueError(f"{name} must use task-state device {device}.")
if pose_batch_size is None:
return value.unsqueeze(0).expand(batch_size, -1, -1).clone()
if pose_batch_size != batch_size:
raise ValueError(
f"{name} batch size must be {batch_size}, got {pose_batch_size}."
)
return value.clone()
def _broadcast_joint_position(
value: torch.Tensor,
*,
batch_size: int,
device: torch.device,
name: str,
) -> torch.Tensor:
"""Resolve an optionally batched joint-position value to a task batch."""
if not isinstance(value, torch.Tensor):
raise TypeError(f"{name} must be a torch.Tensor.")
if value.dim() == 1:
if value.numel() == 0:
raise ValueError(f"{name} must contain at least one joint value.")
value = value.unsqueeze(0).expand(batch_size, -1)
elif value.dim() != 2 or value.shape[0] != batch_size or value.shape[1] == 0:
raise ValueError(
f"{name} must have shape (n_joints,) or " f"({batch_size}, n_joints)."
)
if not value.is_floating_point():
raise TypeError(f"{name} must use a floating-point dtype.")
if value.device != device:
raise ValueError(f"{name} must use task-state device {device}.")
if not torch.isfinite(value).all():
raise ValueError(f"{name} must contain only finite values.")
return value.clone()
[docs]
@dataclass(frozen=True, slots=True, eq=False)
class ArticulationJointState:
"""Verified symbolic state for one named articulation joint."""
position: torch.Tensor
"""Joint positions with shape ``(J,)`` or ``(B, J)``."""
env_mask: torch.Tensor | None = None
"""Rows for which the verified state is present."""
def __post_init__(self) -> None:
if not isinstance(self.position, torch.Tensor):
raise TypeError("ArticulationJointState.position must be a tensor.")
if self.position.dim() not in (1, 2) or self.position.numel() == 0:
raise ValueError(
"ArticulationJointState.position must have shape (J,) or (B, J)."
)
if not self.position.is_floating_point():
raise TypeError("ArticulationJointState.position must be floating point.")
if not torch.isfinite(self.position).all():
raise ValueError("ArticulationJointState.position must be finite.")
object.__setattr__(self, "position", self.position.clone())
if self.env_mask is not None:
batch_size = int(self.position.shape[0]) if self.position.dim() == 2 else -1
if batch_size <= 0:
if self.env_mask.dim() != 1 or self.env_mask.numel() == 0:
raise ValueError(
"ArticulationJointState.env_mask must be a non-empty vector."
)
batch_size = int(self.env_mask.shape[0])
object.__setattr__(
self,
"env_mask",
_normalize_mask(
self.env_mask,
batch_size=batch_size,
device=self.position.device,
name="ArticulationJointState.env_mask",
),
)
[docs]
@dataclass(frozen=True, slots=True, eq=False)
class HeldObjectState:
"""Observed or projected relation between an object and one manipulator."""
semantics: ObjectSemantics
"""Semantics of the held object."""
object_to_eef: torch.Tensor
"""Object-to-end-effector transform."""
grasp_xpos: torch.Tensor
"""End-effector grasp pose."""
env_mask: torch.Tensor | None = None
"""Environments in which the relation is active."""
def __post_init__(self) -> None:
from .core import ObjectSemantics
if not isinstance(self.semantics, ObjectSemantics):
raise TypeError("semantics must be an ObjectSemantics instance.")
object_batch = _validate_pose(self.object_to_eef, "object_to_eef")
grasp_batch = _validate_pose(self.grasp_xpos, "grasp_xpos")
explicit_batches = {
size for size in (object_batch, grasp_batch) if size is not None
}
if len(explicit_batches) > 1:
raise ValueError("Held-object poses must use the same batch size.")
if self.object_to_eef.device != self.grasp_xpos.device:
raise ValueError("Held-object poses must use the same device.")
if self.env_mask is not None:
mask_batch = int(self.env_mask.shape[0]) if self.env_mask.dim() == 1 else -1
batch_size = next(iter(explicit_batches), mask_batch)
object.__setattr__(
self,
"env_mask",
_normalize_mask(
self.env_mask,
batch_size=batch_size,
device=self.object_to_eef.device,
name="env_mask",
),
)
[docs]
@dataclass(frozen=True, slots=True, eq=False)
class CoordinatedHeldObjectState:
"""Observed or projected relation for an object held by two manipulators."""
semantics: ObjectSemantics
left_object_to_eef: torch.Tensor
right_object_to_eef: torch.Tensor
left_grasp_xpos: torch.Tensor
right_grasp_xpos: torch.Tensor
env_mask: torch.Tensor | None = None
def __post_init__(self) -> None:
from .core import ObjectSemantics
if not isinstance(self.semantics, ObjectSemantics):
raise TypeError("semantics must be an ObjectSemantics instance.")
poses = {
"left_object_to_eef": self.left_object_to_eef,
"right_object_to_eef": self.right_object_to_eef,
"left_grasp_xpos": self.left_grasp_xpos,
"right_grasp_xpos": self.right_grasp_xpos,
}
batches = {_validate_pose(value, name) for name, value in poses.items()}
batches.discard(None)
if len(batches) > 1:
raise ValueError("Coordinated held-object poses must share a batch size.")
if len({value.device for value in poses.values()}) != 1:
raise ValueError("Coordinated held-object poses must share a device.")
if self.env_mask is not None:
mask_batch = int(self.env_mask.shape[0]) if self.env_mask.dim() == 1 else -1
batch_size = next(iter(batches), mask_batch)
object.__setattr__(
self,
"env_mask",
_normalize_mask(
self.env_mask,
batch_size=batch_size,
device=self.left_object_to_eef.device,
name="env_mask",
),
)
def _normalize_held(
value: HeldObjectState,
*,
batch_size: int,
device: torch.device,
) -> HeldObjectState:
"""Normalize a held-object relation to one task-state batch."""
return HeldObjectState(
semantics=value.semantics,
object_to_eef=_broadcast_pose(
value.object_to_eef,
batch_size=batch_size,
device=device,
name="HeldObjectState.object_to_eef",
),
grasp_xpos=_broadcast_pose(
value.grasp_xpos,
batch_size=batch_size,
device=device,
name="HeldObjectState.grasp_xpos",
),
env_mask=_normalize_mask(
value.env_mask,
batch_size=batch_size,
device=device,
name="HeldObjectState.env_mask",
),
)
def _normalize_coordinated_held(
value: CoordinatedHeldObjectState,
*,
batch_size: int,
device: torch.device,
) -> CoordinatedHeldObjectState:
"""Normalize a coordinated relation to one task-state batch."""
return CoordinatedHeldObjectState(
semantics=value.semantics,
left_object_to_eef=_broadcast_pose(
value.left_object_to_eef,
batch_size=batch_size,
device=device,
name="CoordinatedHeldObjectState.left_object_to_eef",
),
right_object_to_eef=_broadcast_pose(
value.right_object_to_eef,
batch_size=batch_size,
device=device,
name="CoordinatedHeldObjectState.right_object_to_eef",
),
left_grasp_xpos=_broadcast_pose(
value.left_grasp_xpos,
batch_size=batch_size,
device=device,
name="CoordinatedHeldObjectState.left_grasp_xpos",
),
right_grasp_xpos=_broadcast_pose(
value.right_grasp_xpos,
batch_size=batch_size,
device=device,
name="CoordinatedHeldObjectState.right_grasp_xpos",
),
env_mask=_normalize_mask(
value.env_mask,
batch_size=batch_size,
device=device,
name="CoordinatedHeldObjectState.env_mask",
),
)
def _normalize_articulation_joint(
value: ArticulationJointState,
*,
batch_size: int,
device: torch.device,
) -> ArticulationJointState:
"""Normalize one articulation-joint state to a task-state batch."""
return ArticulationJointState(
position=_broadcast_joint_position(
value.position,
batch_size=batch_size,
device=device,
name="ArticulationJointState.position",
),
env_mask=_normalize_mask(
value.env_mask,
batch_size=batch_size,
device=device,
name="ArticulationJointState.env_mask",
),
)
[docs]
@dataclass(frozen=True, slots=True, eq=False)
class TaskState:
"""Symbolic task state, separate from measured robot state."""
batch_size: int
"""Number of vectorized environments represented by the state."""
device: torch.device | str
"""Device used by per-environment masks and relation tensors."""
held_objects: Mapping[str, HeldObjectState] = field(default_factory=dict)
"""Single-manipulator held-object relations keyed by control resource."""
coordinated_held_objects: Mapping[tuple[str, str], CoordinatedHeldObjectState] = (
field(default_factory=dict)
)
"""Coordinated relations keyed by ordered logical task-state resource pairs."""
articulation_joints: Mapping[tuple[str, str], ArticulationJointState] = field(
default_factory=dict
)
"""Verified articulation states keyed by canonical articulation and joint IDs."""
def __post_init__(self) -> None:
if self.batch_size <= 0:
raise ValueError("TaskState.batch_size must be greater than zero.")
device = _resolve_runtime_device(self.device)
normalized_held: dict[str, HeldObjectState] = {}
for resource, value in self.held_objects.items():
if not isinstance(resource, str) or not resource:
raise TypeError("held_objects keys must be non-empty strings.")
if not isinstance(value, HeldObjectState):
raise TypeError("held_objects values must be HeldObjectState objects.")
normalized_held[resource] = _normalize_held(
value, batch_size=self.batch_size, device=device
)
normalized_coordinated: dict[tuple[str, str], CoordinatedHeldObjectState] = {}
for resources, value in self.coordinated_held_objects.items():
if (
not isinstance(resources, tuple)
or len(resources) != 2
or not all(isinstance(item, str) and item for item in resources)
):
raise TypeError(
"coordinated_held_objects keys must be pairs of non-empty strings."
)
if not isinstance(value, CoordinatedHeldObjectState):
raise TypeError(
"coordinated_held_objects values must be "
"CoordinatedHeldObjectState objects."
)
normalized_coordinated[resources] = _normalize_coordinated_held(
value,
batch_size=self.batch_size,
device=device,
)
normalized_articulation: dict[tuple[str, str], ArticulationJointState] = {}
for key, value in self.articulation_joints.items():
if (
not isinstance(key, tuple)
or len(key) != 2
or not all(
type(item) is str and item and item == item.strip() for item in key
)
):
raise TypeError(
"articulation_joints keys must be pairs of non-empty "
"canonical identifiers."
)
if not isinstance(value, ArticulationJointState):
raise TypeError(
"articulation_joints values must be ArticulationJointState "
"objects."
)
normalized_articulation[key] = _normalize_articulation_joint(
value,
batch_size=self.batch_size,
device=device,
)
object.__setattr__(self, "device", device)
object.__setattr__(self, "held_objects", MappingProxyType(normalized_held))
object.__setattr__(
self,
"coordinated_held_objects",
MappingProxyType(normalized_coordinated),
)
object.__setattr__(
self,
"articulation_joints",
MappingProxyType(normalized_articulation),
)
[docs]
@classmethod
def empty(
cls,
batch_size: int,
device: torch.device | str,
) -> TaskState:
"""Create an empty symbolic state.
Args:
batch_size: Number of represented environments.
device: Tensor device used by the state.
Returns:
Empty task state with explicit batch metadata.
"""
return cls(batch_size=batch_size, device=device)
[docs]
def get_held_object(self, resource: str) -> HeldObjectState | None:
"""Return the object held by ``resource``, if any."""
return self.held_objects.get(resource)
[docs]
def get_coordinated_held_object(
self,
first_resource: str,
second_resource: str,
) -> CoordinatedHeldObjectState | None:
"""Return the relation for an ordered resource pair, if any."""
return self.coordinated_held_objects.get((first_resource, second_resource))
[docs]
def get_articulation_joint_state(
self,
articulation_id: str,
joint_id: str,
) -> ArticulationJointState | None:
"""Return verified state for one canonical articulation joint."""
return self.articulation_joints.get((articulation_id, joint_id))
[docs]
def held_object_mask(self, resource: str) -> torch.Tensor:
"""Return environments where ``resource`` holds an object.
Args:
resource: Manipulator control-resource name.
Returns:
Owned boolean mask with shape ``(batch_size,)``. Missing resources
produce an all-false mask.
"""
held = self.get_held_object(resource)
if held is None:
return torch.zeros(
self.batch_size,
dtype=torch.bool,
device=self.device,
)
assert held.env_mask is not None
return held.env_mask.clone()
[docs]
def exclusive_held_object_mask(self, resource: str) -> torch.Tensor:
"""Return environments where only ``resource`` holds its object.
Object identity is established by the exact semantic record or by a
shared non-null simulation entity. Labels and structural equality are
deliberately ignored because distinct physical objects may look alike.
Args:
resource: Manipulator control-resource name.
Returns:
Owned boolean mask with shape ``(batch_size,)``.
"""
held = self.get_held_object(resource)
if held is None:
return self.held_object_mask(resource)
exclusive = self.held_object_mask(resource)
for other_resource, other in self.held_objects.items():
if other_resource == resource:
continue
if not _same_physical_object(held.semantics, other.semantics):
continue
assert other.env_mask is not None
exclusive &= ~other.env_mask
return exclusive
[docs]
@dataclass(frozen=True, slots=True, eq=False)
class RobotObservation:
"""Measured robot state used as the start of planning or replanning."""
timestamp: float
qpos: torch.Tensor
qvel: torch.Tensor
qeffort: torch.Tensor | None = None
root_pose: torch.Tensor | None = None
root_twist: torch.Tensor | None = None
def __post_init__(self) -> None:
if self.timestamp < 0.0:
raise ValueError("RobotObservation.timestamp must be non-negative.")
if not isinstance(self.qpos, torch.Tensor) or self.qpos.dim() != 2:
raise ValueError(
"RobotObservation.qpos must have shape (num_envs, robot_dof)."
)
if self.qpos.shape[0] == 0 or self.qpos.shape[1] == 0:
raise ValueError("RobotObservation.qpos dimensions must be non-zero.")
if not isinstance(self.qvel, torch.Tensor):
raise TypeError("RobotObservation.qvel must be a torch.Tensor.")
if self.qvel.shape != self.qpos.shape:
raise ValueError("RobotObservation.qvel must match qpos shape.")
if self.qvel.device != self.qpos.device:
raise ValueError("RobotObservation.qpos and qvel must share a device.")
if self.qeffort is not None:
if self.qeffort.shape != self.qpos.shape:
raise ValueError("RobotObservation.qeffort must match qpos shape.")
if self.qeffort.device != self.qpos.device:
raise ValueError("RobotObservation.qeffort must share the qpos device.")
if self.root_pose is not None:
if not isinstance(self.root_pose, torch.Tensor):
raise TypeError("RobotObservation.root_pose must be a tensor or None.")
if self.root_pose.shape != (self.qpos.shape[0], 4, 4):
raise ValueError(
"RobotObservation.root_pose must have shape "
f"({self.qpos.shape[0]}, 4, 4)."
)
if not self.root_pose.is_floating_point():
raise TypeError("RobotObservation.root_pose must be floating point.")
if self.root_pose.device != self.qpos.device:
raise ValueError(
"RobotObservation.root_pose must share the qpos device."
)
if not torch.isfinite(self.root_pose).all():
raise ValueError("RobotObservation.root_pose must be finite.")
if self.root_twist is not None:
if not isinstance(self.root_twist, torch.Tensor):
raise TypeError("RobotObservation.root_twist must be a tensor or None.")
if self.root_twist.shape != (self.qpos.shape[0], 6):
raise ValueError(
"RobotObservation.root_twist must have shape "
f"({self.qpos.shape[0]}, 6)."
)
if not self.root_twist.is_floating_point():
raise TypeError("RobotObservation.root_twist must be floating point.")
if self.root_twist.device != self.qpos.device:
raise ValueError(
"RobotObservation.root_twist must share the qpos device."
)
if not torch.isfinite(self.root_twist).all():
raise ValueError("RobotObservation.root_twist must be finite.")
object.__setattr__(self, "qpos", self.qpos.clone())
object.__setattr__(self, "qvel", self.qvel.clone())
if self.qeffort is not None:
object.__setattr__(self, "qeffort", self.qeffort.clone())
if self.root_pose is not None:
object.__setattr__(self, "root_pose", self.root_pose.clone())
if self.root_twist is not None:
object.__setattr__(self, "root_twist", self.root_twist.clone())
@property
def batch_size(self) -> int:
"""Number of represented vectorized environments."""
return int(self.qpos.shape[0])
@property
def robot_dof(self) -> int:
"""Number of robot joint-position columns."""
return int(self.qpos.shape[1])
[docs]
def with_qpos(self, qpos: torch.Tensor) -> RobotObservation:
"""Create a projected observation with a new position and zero velocity.
Args:
qpos: Projected joint positions with the same shape as this observation.
Returns:
New observation suitable for compiling the next action.
"""
return RobotObservation(
timestamp=self.timestamp,
qpos=qpos,
qvel=torch.zeros_like(qpos),
qeffort=self.qeffort,
root_pose=self.root_pose,
root_twist=self.root_twist,
)
[docs]
@dataclass(frozen=True, slots=True, eq=False)
class EntityState:
"""Scene entity state addressable by a stable entity identifier."""
pose: torch.Tensor
confidence: float = 1.0
def __post_init__(self) -> None:
_validate_pose(self.pose, "EntityState.pose")
if not 0.0 <= self.confidence <= 1.0:
raise ValueError("EntityState.confidence must be in [0, 1].")
object.__setattr__(self, "pose", self.pose.clone())
[docs]
@dataclass(frozen=True, slots=True, eq=False)
class ObservedArticulationJointState:
"""Live measured state for one scene articulation joint."""
position: torch.Tensor
"""Measured joint position with shape ``(J,)`` or ``(B, J)``."""
valid_mask: torch.Tensor | None = None
"""Optional row-validity mask for a batched observation."""
def __post_init__(self) -> None:
position = self.position
if not isinstance(position, torch.Tensor):
raise TypeError("ObservedArticulationJointState.position must be a tensor.")
if position.dim() not in (1, 2) or position.numel() == 0:
raise ValueError(
"ObservedArticulationJointState.position must have shape (J,) "
"or (B, J)."
)
if not position.is_floating_point():
raise TypeError(
"ObservedArticulationJointState.position must be floating point."
)
if not torch.isfinite(position).all():
raise ValueError(
"ObservedArticulationJointState.position must contain only "
"finite values."
)
object.__setattr__(self, "position", position.clone())
if self.valid_mask is None:
return
valid_mask = self.valid_mask
if not isinstance(valid_mask, torch.Tensor):
raise TypeError(
"ObservedArticulationJointState.valid_mask must be a tensor or None."
)
if position.dim() != 2:
raise ValueError(
"ObservedArticulationJointState.valid_mask requires a batched "
"position."
)
if valid_mask.dtype != torch.bool or valid_mask.shape != (position.shape[0],):
raise ValueError(
"ObservedArticulationJointState.valid_mask must have shape (B,) "
"and dtype torch.bool."
)
if valid_mask.device != position.device:
raise ValueError(
"ObservedArticulationJointState position and valid_mask must "
"share a device."
)
object.__setattr__(self, "valid_mask", valid_mask.clone())
[docs]
def snapshot(self) -> ObservedArticulationJointState:
"""Return an independently owned observation value."""
return ObservedArticulationJointState(self.position, self.valid_mask)
class _ImmutableEntityMapping(Mapping[str, EntityState]):
"""Own entity states and return defensive copies on every public read."""
__slots__ = ("_states",)
def __init__(self, states: Mapping[str, EntityState]) -> None:
self._states = MappingProxyType(
{
entity_id: EntityState(state.pose, confidence=state.confidence)
for entity_id, state in states.items()
}
)
def __getitem__(self, entity_id: str) -> EntityState:
state = self._states[entity_id]
return EntityState(state.pose, confidence=state.confidence)
def __iter__(self) -> Iterator[str]:
return iter(self._states)
def __len__(self) -> int:
return len(self._states)
class _ImmutableObservedArticulationJointMapping(
Mapping[tuple[str, str], ObservedArticulationJointState]
):
"""Own live joint observations and copy values on every public read."""
__slots__ = ("_states",)
def __init__(
self,
states: Mapping[tuple[str, str], ObservedArticulationJointState],
) -> None:
self._states = MappingProxyType(
{key: state.snapshot() for key, state in states.items()}
)
def __getitem__(
self,
key: tuple[str, str],
) -> ObservedArticulationJointState:
return self._states[key].snapshot()
def __iter__(self) -> Iterator[tuple[str, str]]:
return iter(self._states)
def __len__(self) -> int:
return len(self._states)
[docs]
@dataclass(frozen=True, slots=True, eq=False)
class SceneSnapshot:
"""Versioned scene state used to ground dynamic goals and obstacles."""
timestamp: float
version: int
entities: Mapping[str, EntityState] = field(default_factory=dict)
collision_world_revision: int | tuple[int, ...] = 0
"""Global or per-environment collision-world revision."""
collision_entity_ids: tuple[str, ...] = ()
"""Entity IDs whose poses update a planner's dynamic collision world."""
articulation_joints: Mapping[tuple[str, str], ObservedArticulationJointState] = (
field(default_factory=dict)
)
"""Live physical joint observations keyed by articulation and joint ID."""
def __post_init__(self) -> None:
if self.timestamp < 0.0:
raise ValueError("SceneSnapshot.timestamp must be non-negative.")
if self.version < 0:
raise ValueError("SceneSnapshot.version must be non-negative.")
revision = self.collision_world_revision
if isinstance(revision, bool):
raise TypeError("collision_world_revision must contain integers.")
if isinstance(revision, int):
if revision < 0:
raise ValueError(
"collision_world_revision must contain non-negative values."
)
else:
if not isinstance(revision, tuple) or not revision:
raise TypeError(
"collision_world_revision must be an integer or a non-empty "
"tuple of integers."
)
if any(
isinstance(value, bool) or not isinstance(value, int)
for value in revision
):
raise TypeError("collision_world_revision must contain integers.")
if any(value < 0 for value in revision):
raise ValueError(
"collision_world_revision must contain non-negative values."
)
normalized: dict[str, EntityState] = {}
for entity_id, state in self.entities.items():
if not isinstance(entity_id, str) or not entity_id:
raise ValueError("Scene entity identifiers must be non-empty strings.")
if not isinstance(state, EntityState):
raise TypeError(
"SceneSnapshot entities must contain EntityState values."
)
normalized[entity_id] = state
normalized_joints: dict[tuple[str, str], ObservedArticulationJointState] = {}
for key, state in self.articulation_joints.items():
if (
not isinstance(key, tuple)
or len(key) != 2
or not all(
type(identifier) is str
and identifier
and identifier == identifier.strip()
for identifier in key
)
):
raise TypeError(
"SceneSnapshot articulation_joints keys must be canonical "
"(articulation_id, joint_id) pairs."
)
if not isinstance(state, ObservedArticulationJointState):
raise TypeError(
"SceneSnapshot articulation_joints values must be "
"ObservedArticulationJointState objects."
)
normalized_joints[key] = state
collision_entity_ids = tuple(self.collision_entity_ids)
if len(set(collision_entity_ids)) != len(collision_entity_ids) or not all(
isinstance(entity_id, str) and entity_id
for entity_id in collision_entity_ids
):
raise ValueError(
"collision_entity_ids must contain unique non-empty entity IDs."
)
missing = set(collision_entity_ids).difference(normalized)
if missing:
raise ValueError(
"collision_entity_ids reference missing scene entities: "
f"{sorted(missing)}."
)
object.__setattr__(self, "entities", _ImmutableEntityMapping(normalized))
object.__setattr__(
self,
"articulation_joints",
_ImmutableObservedArticulationJointMapping(normalized_joints),
)
object.__setattr__(self, "collision_entity_ids", collision_entity_ids)
[docs]
def get_articulation_joint_state(
self,
articulation_id: str,
joint_id: str,
) -> ObservedArticulationJointState | None:
"""Return an owned live joint observation for a canonical address."""
for value, field_name in (
(articulation_id, "articulation_id"),
(joint_id, "joint_id"),
):
if type(value) is not str or not value or value != value.strip():
raise ValueError(
f"{field_name} must be a non-empty canonical identifier."
)
return self.articulation_joints.get((articulation_id, joint_id))
[docs]
def collision_world_revisions(self, batch_size: int) -> tuple[int, ...]:
"""Expand the collision revision to one value per environment.
Args:
batch_size: Number of environments represented by the planning context.
Returns:
Per-environment monotonic revision tuple.
Raises:
ValueError: If an explicit revision tuple does not match the batch.
"""
if batch_size <= 0:
raise ValueError("batch_size must be positive.")
revision = self.collision_world_revision
if isinstance(revision, int):
return (revision,) * batch_size
if len(revision) == 1:
return revision * batch_size
if len(revision) != batch_size:
raise ValueError(
"collision_world_revision must be global or have one value per "
f"environment; got {len(revision)} values for batch {batch_size}."
)
return revision
[docs]
def collision_obstacle_poses(
self,
*,
batch_size: int,
device: torch.device,
dtype: torch.dtype,
) -> Mapping[str, torch.Tensor]:
"""Return collision obstacle poses in planning batch order.
Args:
batch_size: Number of planning environments.
device: Planner tensor device.
dtype: Planner tensor dtype.
Returns:
Mapping from configured collision entity ID to ``(B, 4, 4)`` pose.
"""
poses: dict[str, torch.Tensor] = {}
for entity_id in self.collision_entity_ids:
pose = self.entities[entity_id].pose.to(device=device, dtype=dtype)
if pose.shape == (4, 4):
pose = pose.unsqueeze(0).expand(batch_size, -1, -1)
elif pose.shape != (batch_size, 4, 4):
raise ValueError(
f"Collision entity {entity_id!r} pose must match planning "
f"batch size {batch_size}."
)
poses[entity_id] = pose.clone()
return MappingProxyType(poses)
[docs]
@classmethod
def empty(cls) -> SceneSnapshot:
"""Create an empty initial scene snapshot."""
return cls(timestamp=0.0, version=0)
[docs]
@dataclass(frozen=True, slots=True, eq=False)
class PlanningContext:
"""Complete side-effect-free input to :meth:`AtomicAction.plan`."""
robot: RobotObservation
task: TaskState
scene: SceneSnapshot
env_ids: torch.Tensor
control_dt: float | None = None
"""Explicit command period used by action-owned interpolation."""
def __post_init__(self) -> None:
if not isinstance(self.robot, RobotObservation):
raise TypeError("robot must be a RobotObservation.")
if not isinstance(self.task, TaskState):
raise TypeError("task must be a TaskState.")
if not isinstance(self.scene, SceneSnapshot):
raise TypeError("scene must be a SceneSnapshot.")
if self.task.batch_size != self.robot.batch_size:
raise ValueError("TaskState and RobotObservation batch sizes must match.")
if self.task.device != self.robot.qpos.device:
raise ValueError("TaskState and RobotObservation must share a device.")
self.scene.collision_world_revisions(self.robot.batch_size)
for entity_id, state in self.scene.entities.items():
if state.pose.dim() == 3 and state.pose.shape[0] != self.robot.batch_size:
raise ValueError(
f"Scene entity {entity_id!r} pose batch must match the "
"planning context."
)
for (
articulation_id,
joint_id,
), state in self.scene.articulation_joints.items():
if (
state.position.dim() == 2
and state.position.shape[0] != self.robot.batch_size
):
raise ValueError(
f"Scene articulation joint ({articulation_id!r}, {joint_id!r}) "
"position batch must match the planning context."
)
if not isinstance(self.env_ids, torch.Tensor):
raise TypeError("env_ids must be a torch.Tensor.")
if self.env_ids.dtype != torch.long:
raise TypeError("env_ids must have dtype torch.long.")
if self.env_ids.shape != (self.robot.batch_size,):
raise ValueError(
"env_ids must identify every row in the planning batch; expected "
f"shape ({self.robot.batch_size},), got {tuple(self.env_ids.shape)}."
)
if self.env_ids.device != self.robot.qpos.device:
raise ValueError("env_ids and robot tensors must share a device.")
if torch.unique(self.env_ids).numel() != self.env_ids.numel():
raise ValueError("env_ids must be unique.")
if self.control_dt is not None:
if isinstance(self.control_dt, bool) or not isinstance(
self.control_dt, (int, float)
):
raise TypeError("control_dt must be a real number or None.")
if not math.isfinite(self.control_dt) or self.control_dt <= 0.0:
raise ValueError("control_dt must be finite and greater than zero.")
object.__setattr__(self, "control_dt", float(self.control_dt))
object.__setattr__(self, "env_ids", self.env_ids.clone())
@property
def batch_size(self) -> int:
"""Number of environments in this planning request."""
return self.robot.batch_size
@property
def last_qpos(self) -> torch.Tensor:
"""Measured joint positions used as the planning start state."""
return self.robot.qpos
@property
def held_objects(self) -> Mapping[str, HeldObjectState]:
"""Single-resource held-object relations."""
return self.task.held_objects
@property
def coordinated_held_objects(
self,
) -> Mapping[tuple[str, str], CoordinatedHeldObjectState]:
"""Coordinated held-object relations."""
return self.task.coordinated_held_objects
[docs]
def get_held_object(self, resource: str) -> HeldObjectState | None:
"""Return the object held by ``resource``, if any."""
return self.task.get_held_object(resource)
[docs]
def get_coordinated_held_object(
self,
first_resource: str,
second_resource: str,
) -> CoordinatedHeldObjectState | None:
"""Return a coordinated held-object relation, if any."""
return self.task.get_coordinated_held_object(first_resource, second_resource)
@property
def articulation_joints(
self,
) -> Mapping[tuple[str, str], ArticulationJointState]:
"""Verified articulation-joint states."""
return self.task.articulation_joints
[docs]
def get_articulation_joint_state(
self,
articulation_id: str,
joint_id: str,
) -> ArticulationJointState | None:
"""Return verified state for one canonical articulation joint."""
return self.task.get_articulation_joint_state(articulation_id, joint_id)
[docs]
def require_control_dt(self) -> float:
"""Return the explicit command period required for interpolation.
Raises:
ValueError: If the caller did not provide ``control_dt``.
"""
if self.control_dt is None:
raise ValueError(
"This action performs interpolation and requires an explicit "
"PlanningContext.control_dt."
)
return self.control_dt
[docs]
def project(
self,
*,
qpos: torch.Tensor,
task: TaskState,
) -> PlanningContext:
"""Create the hypothetical context used to compile a following action.
Args:
qpos: Projected terminal joint positions.
task: Task state after applying expected effects.
Returns:
New context. No measured state or simulator state is mutated.
"""
return PlanningContext(
robot=self.robot.with_qpos(qpos),
task=task,
scene=self.scene,
env_ids=self.env_ids,
control_dt=self.control_dt,
)
__all__ = [
"ArticulationJointState",
"CoordinatedHeldObjectState",
"EntityState",
"HeldObjectState",
"ObservedArticulationJointState",
"PlanningContext",
"RobotObservation",
"SceneSnapshot",
"TaskState",
]