# ----------------------------------------------------------------------------
# Copyright (c) 2021-2026 DexForce Technology Co., Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ----------------------------------------------------------------------------
from __future__ import annotations
import hashlib
import uuid
from dataclasses import dataclass
from time import perf_counter
from typing import TYPE_CHECKING
from urllib.parse import quote
import numpy as np
from ._utils import to_numpy_array as _to_numpy
from .cfg import VisualizationCfg
from .protocol import (
CameraImage,
CameraImageFrame,
CameraSpec,
DynamicMeshUpdate,
FrameOverlay,
GizmoSpec,
GizmoState,
JointControlProvider,
JointControlSpec,
JointControlState,
MeshGeometry,
PointCloudOverlay,
SceneFrame,
SceneManifest,
SceneNode,
SceneOverlays,
pose_to_position_wxyz,
)
if TYPE_CHECKING:
from embodichain.lab.sim import SimulationManager
__all__ = [
"CameraImageCaptureResult",
"CaptureResult",
"SceneExporter",
"mesh_geometry_id",
"safe_path_component",
]
def safe_path_component(value: object) -> str:
"""Encode an arbitrary identifier as one stable Viser path component."""
encoded = quote(str(value), safe="-_.~")
return encoded or "unnamed"
def mesh_geometry_id(
vertices: object,
faces: object,
color: tuple[int, int, int] | None = None,
) -> str:
"""Compute a stable content identifier for local mesh geometry and color."""
vertices_array = _to_numpy(vertices, np.float32, copy=False)
faces_array = _to_numpy(faces, np.uint32, copy=False)
digest = hashlib.sha256()
digest.update(str(vertices_array.shape).encode("ascii"))
digest.update(vertices_array.tobytes())
digest.update(str(faces_array.shape).encode("ascii"))
digest.update(faces_array.tobytes())
if color is not None:
digest.update(b"color:")
digest.update(bytes(color))
return f"sha256:{digest.hexdigest()}"
[docs]
@dataclass(frozen=True)
class CaptureResult:
"""A captured frame and the time spent copying scene data."""
frame: SceneFrame
capture_seconds: float
[docs]
@dataclass(frozen=True)
class CameraImageCaptureResult:
"""Captured RGB images and the time spent rendering and copying them."""
frame: CameraImageFrame
capture_seconds: float
@dataclass(frozen=True)
class _NodeSource:
node: SceneNode
asset_key: tuple[str, str]
asset: object
link_index: int | None = None
object_index: int | None = None
@dataclass(frozen=True)
class _PoseBatch:
"""Vectorized pose selection for all exported nodes of one asset."""
asset_kind: str
asset: object
destination_indices: np.ndarray
env_ids: np.ndarray
component_indices: np.ndarray | None
@dataclass(frozen=True)
class _CameraSource:
spec: CameraSpec
camera: object
@dataclass(frozen=True)
class _GizmoSource:
spec: GizmoSpec
gizmo: object
[docs]
class SceneExporter:
"""Capture backend-neutral scene data from a ``SimulationManager``.
Scene access happens only in :meth:`build_manifest` and :meth:`capture`, so
callers can keep both methods on the simulation thread. Returned protocol
objects own detached CPU copies and are safe to hand to a background backend.
Args:
sim: Simulation manager whose public asset interfaces are exported.
cfg: Visualization selection and sampling configuration.
run_id: Stable identifier for this runtime. A UUID is generated by default.
"""
_COLORS = {
"rigid_object": (90, 200, 255),
"rigid_group_object": (125, 185, 255),
"robot_link": (80, 210, 140),
"articulation_link": (255, 175, 70),
"soft_object": (230, 120, 255),
"cloth_object": (255, 105, 145),
}
[docs]
def __init__(
self,
sim: SimulationManager,
cfg: VisualizationCfg,
run_id: str | None = None,
) -> None:
self._sim = sim
self.cfg = cfg
self.run_id = run_id or str(uuid.uuid4())
self._scene_revision = 0
self._sequence = 0
self._image_sequence = 0
self._sources: tuple[_NodeSource, ...] = ()
self._pose_batches: tuple[_PoseBatch, ...] = ()
self._source_visibility = np.empty((0,), dtype=np.bool_)
self._dynamic_source_indices = np.empty((0,), dtype=np.int64)
self._dynamic_env_ids = np.empty((0,), dtype=np.int64)
self._camera_sources: tuple[_CameraSource, ...] = ()
self._gizmo_sources: tuple[_GizmoSource, ...] = ()
self._joint_control_provider: JointControlProvider | None = None
self._joint_control_specs: tuple[JointControlSpec, ...] = ()
self._env_ids = (
tuple(range(sim.num_envs)) if cfg.env_ids is None else tuple(cfg.env_ids)
)
self._env_offsets = np.zeros((sim.num_envs, 3), dtype=np.float32)
self._validate_env_ids()
@property
def scene_revision(self) -> int:
"""Current manifest revision, or zero before the first manifest."""
return self._scene_revision
@property
def has_cameras(self) -> bool:
"""Whether the current manifest contains an RGB-capable camera."""
return any(
bool(getattr(source.camera.cfg, "enable_color", False))
for source in self._camera_sources
)
@property
def has_deformables(self) -> bool:
"""Whether the current manifest contains soft-body or cloth nodes."""
return any(source.node.dynamic_geometry for source in self._sources)
[docs]
def set_joint_control_provider(
self,
provider: JointControlProvider | None,
) -> None:
"""Install an optional simulation-thread joint-control source.
The next :meth:`build_manifest` call snapshots the provider's static
controls. Dynamic values are then sampled by :meth:`capture`.
"""
self._joint_control_provider = provider
def _validate_env_ids(self) -> None:
invalid = [env_id for env_id in self._env_ids if env_id >= self._sim.num_envs]
if invalid:
raise ValueError(
f"Visualization env_ids {invalid} are outside simulation range "
f"[0, {self._sim.num_envs - 1}]."
)
if (
self.cfg.max_visible_envs is not None
and len(self._env_ids) > self.cfg.max_visible_envs
):
raise ValueError(
f"Selected {len(self._env_ids)} environments, exceeding "
f"max_visible_envs={self.cfg.max_visible_envs}."
)
@staticmethod
def _build_pose_batches(sources: tuple[_NodeSource, ...]) -> tuple[_PoseBatch, ...]:
"""Compile per-asset node selections into NumPy indexing arrays."""
grouped: dict[tuple[str, str], list[tuple[int, _NodeSource]]] = {}
for destination_index, source in enumerate(sources):
if source.node.dynamic_geometry:
continue
grouped.setdefault(source.asset_key, []).append((destination_index, source))
batches: list[_PoseBatch] = []
for asset_key, indexed_sources in grouped.items():
component_indices = [
(
source.link_index
if source.link_index is not None
else source.object_index
)
for _, source in indexed_sources
]
has_components = component_indices[0] is not None
if any(
(index is not None) != has_components for index in component_indices
):
raise ValueError(
f"Visualization asset {asset_key!r} mixes component and root poses."
)
batches.append(
_PoseBatch(
asset_kind=asset_key[0],
asset=indexed_sources[0][1].asset,
destination_indices=np.asarray(
[index for index, _ in indexed_sources],
dtype=np.int64,
),
env_ids=np.asarray(
[source.node.env_id for _, source in indexed_sources],
dtype=np.int64,
),
component_indices=(
np.asarray(component_indices, dtype=np.int64)
if has_components
else None
),
)
)
return tuple(batches)
def _read_arena_offsets(self) -> np.ndarray:
offsets = _to_numpy(self._sim.arena_offsets, np.float32)
if offsets.shape != (self._sim.num_envs, 3):
raise ValueError(
"SimulationManager.arena_offsets must have shape "
f"({self._sim.num_envs}, 3), received {offsets.shape}."
)
return offsets
@staticmethod
def _add_geometry(
geometries: dict[str, MeshGeometry],
vertices: object,
faces: object,
color: tuple[int, int, int],
) -> str | None:
vertices_array = _to_numpy(vertices, np.float32, copy=False)
faces_array = _to_numpy(faces, np.uint32, copy=False)
if vertices_array.size == 0 or faces_array.size == 0:
return None
geometry_id = mesh_geometry_id(vertices_array, faces_array, color)
if geometry_id not in geometries:
geometries[geometry_id] = MeshGeometry(
geometry_id=geometry_id,
vertices=vertices_array,
faces=faces_array,
color=color,
)
return geometry_id
[docs]
def build_manifest(self) -> SceneManifest:
"""Capture scene topology, incrementing the scene revision.
Call this after assets are created and again after topology-changing
resets. Empty articulation links are ignored.
"""
self._validate_env_ids()
self._env_offsets = self._read_arena_offsets()
geometries: dict[str, MeshGeometry] = {}
sources: list[_NodeSource] = []
camera_sources: list[_CameraSource] = []
gizmo_sources: list[_GizmoSource] = []
for uid in self._sim.get_rigid_object_uid_list():
asset = self._sim.get_rigid_object(uid)
if asset is None:
continue
selected_env_ids = list(self._env_ids)
vertices_by_env = asset.get_vertices(
env_ids=selected_env_ids,
scale=True,
)
faces_by_env = asset.get_triangles(env_ids=selected_env_ids)
for selected_index, env_id in enumerate(self._env_ids):
vertices = vertices_by_env[selected_index]
faces = faces_by_env[selected_index]
geometry_id = self._add_geometry(
geometries, vertices, faces, self._COLORS["rigid_object"]
)
if geometry_id is None:
continue
uid_component = safe_path_component(uid)
node = SceneNode(
node_id=f"env:{env_id}/rigid:{uid_component}",
path=f"/envs/{env_id}/rigid_objects/{uid_component}",
parent_id=f"env:{env_id}",
env_id=env_id,
kind="rigid_object",
geometry_id=geometry_id,
)
sources.append(
_NodeSource(node=node, asset_key=("rigid", uid), asset=asset)
)
self._append_rigid_object_groups(sources, geometries)
self._append_articulations(
sources=sources,
geometries=geometries,
uids=self._sim.get_robot_uid_list(),
getter=self._sim.get_robot,
kind="robot_link",
asset_prefix="robot",
)
self._append_articulations(
sources=sources,
geometries=geometries,
uids=self._sim.get_articulation_uid_list(),
getter=self._sim.get_articulation,
kind="articulation_link",
asset_prefix="articulation",
)
self._append_deformable_objects(
sources=sources,
geometries=geometries,
uids=self._sim.get_soft_object_uid_list(),
getter=self._sim.get_soft_object,
kind="soft_object",
asset_prefix="soft",
)
self._append_deformable_objects(
sources=sources,
geometries=geometries,
uids=self._sim.get_cloth_object_uid_list(),
getter=self._sim.get_cloth_object,
kind="cloth_object",
asset_prefix="cloth",
)
self._append_cameras(camera_sources)
self._append_gizmos(gizmo_sources)
joint_control_specs: tuple[JointControlSpec, ...] = ()
if self._joint_control_provider is not None:
joint_control_specs = tuple(
spec
for spec in self._joint_control_provider.joint_control_specs()
if spec.env_id in self._env_ids
)
self._scene_revision += 1
self._sources = tuple(sources)
self._pose_batches = self._build_pose_batches(self._sources)
self._source_visibility = np.asarray(
[source.node.visible for source in self._sources],
dtype=np.bool_,
)
self._dynamic_source_indices = np.asarray(
[
index
for index, source in enumerate(self._sources)
if source.node.dynamic_geometry
],
dtype=np.int64,
)
self._dynamic_env_ids = np.asarray(
[
self._sources[index].node.env_id
for index in self._dynamic_source_indices
],
dtype=np.int64,
)
self._camera_sources = tuple(camera_sources)
self._gizmo_sources = tuple(gizmo_sources)
self._joint_control_specs = joint_control_specs
return SceneManifest(
run_id=self.run_id,
scene_revision=self._scene_revision,
nodes=tuple(source.node for source in sources),
geometries=tuple(geometries.values()),
cameras=tuple(source.spec for source in camera_sources),
gizmos=tuple(source.spec for source in gizmo_sources),
joint_controls=joint_control_specs,
)
def _capture_joint_control_states(self) -> tuple[JointControlState, ...]:
if self._joint_control_provider is None or not self._joint_control_specs:
return ()
states = {
state.control_id: state
for state in self._joint_control_provider.joint_control_states()
}
missing = [
spec.control_id
for spec in self._joint_control_specs
if spec.control_id not in states
]
if missing:
raise ValueError(
f"Joint control provider omitted states for controls: {missing}."
)
return tuple(states[spec.control_id] for spec in self._joint_control_specs)
def _append_gizmos(self, sources: list[_GizmoSource]) -> None:
if 0 not in self._env_ids:
return
get_gizmo_items = getattr(self._sim, "get_gizmo_items", None)
if get_gizmo_items is None:
return
for gizmo_id, gizmo in get_gizmo_items():
target = gizmo.target
if target is None:
continue
target_uid = str(
getattr(getattr(target, "cfg", None), "uid", None) or gizmo_id
)
scale = max(
float(gizmo.cfg.axis_length_x),
float(gizmo.cfg.axis_length_y),
float(gizmo.cfg.axis_length_z),
float(gizmo.cfg.rings_radius),
)
line_width = max(
1.0,
float(max(gizmo.cfg.axis_size, gizmo.cfg.rings_size)) * 250.0,
)
sources.append(
_GizmoSource(
spec=GizmoSpec(
gizmo_id=gizmo_id,
target_uid=target_uid,
target_type=gizmo.target_type,
control_part=gizmo.control_part,
env_id=0,
path=f"/interactions/gizmos/{safe_path_component(gizmo_id)}",
scale=scale,
line_width=line_width,
visible=gizmo.is_visible(),
),
gizmo=gizmo,
)
)
[docs]
def resolve_node_target(self, node_id: str) -> tuple[str, str] | None:
"""Map a published scene node id to its ``(uid, kind)``.
Used by the simulation thread to turn a Viser click-pick result into the
asset uid that :meth:`SimulationManager.enable_gizmo` expects.
Args:
node_id: Scene node id from the current manifest.
Returns:
``(uid, kind)`` where ``kind`` is the asset kind (for example
``"rigid"``, ``"robot"``, or ``"articulation"``), or ``None`` if the
node id is not part of the current scene.
"""
for source in self._sources:
if source.node.node_id == node_id:
kind, uid = source.asset_key
return str(uid), str(kind)
return None
def _append_rigid_object_groups(
self,
sources: list[_NodeSource],
geometries: dict[str, MeshGeometry],
) -> None:
for uid in self._sim.get_rigid_object_group_uid_list():
asset = self._sim.get_rigid_object_group(uid)
if asset is None:
continue
uid_component = safe_path_component(uid)
object_names = tuple(asset.cfg.rigid_objects)
selected_env_ids = list(self._env_ids)
for object_index, object_name in enumerate(object_names):
object_component = safe_path_component(object_name)
vertices_by_env = asset.get_object_vertices(
object_index,
env_ids=selected_env_ids,
scale=True,
)
faces_by_env = asset.get_object_triangles(
object_index,
env_ids=selected_env_ids,
)
for selected_index, env_id in enumerate(self._env_ids):
vertices = vertices_by_env[selected_index]
faces = faces_by_env[selected_index]
geometry_id = self._add_geometry(
geometries,
vertices,
faces,
self._COLORS["rigid_group_object"],
)
if geometry_id is None:
continue
node = SceneNode(
node_id=(
f"env:{env_id}/rigid_group:{uid_component}/"
f"object:{object_component}"
),
path=(
f"/envs/{env_id}/rigid_object_groups/{uid_component}/"
f"objects/{object_component}"
),
parent_id=f"env:{env_id}/rigid_group:{uid_component}",
env_id=env_id,
kind="rigid_group_object",
geometry_id=geometry_id,
)
sources.append(
_NodeSource(
node=node,
asset_key=("rigid_group", uid),
asset=asset,
object_index=object_index,
)
)
def _append_deformable_objects(
self,
*,
sources: list[_NodeSource],
geometries: dict[str, MeshGeometry],
uids: list[str],
getter: object,
kind: str,
asset_prefix: str,
) -> None:
for uid in uids:
asset = getter(uid)
if asset is None:
continue
if kind == "soft_object":
current_vertices = _to_numpy(
asset.get_current_collision_vertices(),
np.float32,
)
else:
current_vertices = _to_numpy(
asset.get_current_vertex_position(),
np.float32,
)
uid_component = safe_path_component(uid)
selected_env_ids = list(self._env_ids)
if kind == "soft_object":
faces_by_env = asset.get_collision_surface_triangles(
env_ids=selected_env_ids,
)
else:
faces_by_env = asset.get_triangles(env_ids=selected_env_ids)
for selected_index, env_id in enumerate(self._env_ids):
vertices = current_vertices[env_id] - self._env_offsets[env_id]
faces = faces_by_env[selected_index]
geometry_id = self._add_geometry(
geometries,
vertices,
faces,
self._COLORS[kind],
)
if geometry_id is None:
continue
node = SceneNode(
node_id=f"env:{env_id}/{asset_prefix}:{uid_component}",
path=f"/envs/{env_id}/{asset_prefix}_objects/{uid_component}",
parent_id=f"env:{env_id}",
env_id=env_id,
kind=kind,
geometry_id=geometry_id,
dynamic_geometry=True,
)
sources.append(
_NodeSource(
node=node,
asset_key=(asset_prefix, uid),
asset=asset,
)
)
@staticmethod
def _camera_fov_y(intrinsics: np.ndarray, height: int) -> float:
if intrinsics.shape == (3, 3):
fy = float(intrinsics[1, 1])
elif intrinsics.shape == (4,):
fy = float(intrinsics[1])
else:
raise ValueError(
"Camera intrinsics must have shape (3, 3) or (4,), "
f"received {intrinsics.shape}."
)
if fy <= 0.0:
raise ValueError("Camera fy must be greater than zero.")
return float(2.0 * np.arctan(float(height) / (2.0 * fy)))
def _append_cameras(self, sources: list[_CameraSource]) -> None:
for uid in self._sim.get_sensor_uid_list():
camera = self._sim.get_sensor(uid)
sensor_type = getattr(getattr(camera, "cfg", None), "sensor_type", None)
if camera is None or sensor_type not in {"Camera", "StereoCamera"}:
continue
camera_intrinsics = camera.get_intrinsics()
if sensor_type == "StereoCamera":
# A stereo sensor is represented by its primary (left) RGB
# observation. This keeps one preview per configured sensor,
# matching the environment's ``color`` observation key.
camera_intrinsics = camera_intrinsics[0]
intrinsics = _to_numpy(camera_intrinsics, np.float32)
if intrinsics.shape[0] != self._sim.num_envs:
raise ValueError(
f"Camera {uid!r} returned {intrinsics.shape[0]} intrinsics "
f"for {self._sim.num_envs} environments."
)
uid_component = safe_path_component(uid)
width = int(camera.cfg.width)
height = int(camera.cfg.height)
for env_id in self._env_ids:
camera_id = f"env:{env_id}/camera:{uid_component}"
sources.append(
_CameraSource(
spec=CameraSpec(
camera_id=camera_id,
sensor_uid=uid,
env_id=env_id,
path=f"/envs/{env_id}/cameras/{uid_component}",
fov_y=self._camera_fov_y(intrinsics[env_id], height),
aspect=float(width) / float(height),
near=float(camera.cfg.near),
far=float(camera.cfg.far),
role=getattr(camera.cfg, "visualization_role", "sensor"),
),
camera=camera,
)
)
def _append_articulations(
self,
*,
sources: list[_NodeSource],
geometries: dict[str, MeshGeometry],
uids: list[str],
getter: object,
kind: str,
asset_prefix: str,
) -> None:
for uid in uids:
asset = getter(uid)
if asset is None:
continue
uid_component = safe_path_component(uid)
for link_index, link_name in enumerate(asset.link_names):
vertices, faces = asset.get_link_vert_face(link_name)
geometry_id = self._add_geometry(
geometries, vertices, faces, self._COLORS[kind]
)
if geometry_id is None:
continue
link_component = safe_path_component(link_name)
for env_id in self._env_ids:
node = SceneNode(
node_id=(
f"env:{env_id}/{asset_prefix}:{uid_component}/"
f"link:{link_component}"
),
path=(
f"/envs/{env_id}/{asset_prefix}s/{uid_component}/"
f"links/{link_component}"
),
parent_id=f"env:{env_id}/{asset_prefix}:{uid_component}",
env_id=env_id,
kind=kind,
geometry_id=geometry_id,
)
sources.append(
_NodeSource(
node=node,
asset_key=(asset_prefix, uid),
asset=asset,
link_index=link_index,
)
)
def _capture_axis_marker_overlays(
self, reserved_frame_ids: set[str] | None = None
) -> tuple[FrameOverlay, ...]:
"""Capture native simulation axes as Viser coordinate-frame overlays.
Args:
reserved_frame_ids: Caller-owned frame IDs that generated markers
must not replace.
"""
get_axis_marker_items = getattr(self._sim, "get_axis_marker_items", None)
if get_axis_marker_items is None:
return ()
frames: list[FrameOverlay] = []
used_frame_ids = set(reserved_frame_ids or ())
for marker_name, handles, axis_length, axis_radius in get_axis_marker_items():
for index, handle in enumerate(handles):
position, wxyz = pose_to_position_wxyz(handle.get_world_pose())
base_id = f"marker:{marker_name}:{index}"
overlay_id = base_id
suffix = 1
while overlay_id in used_frame_ids:
overlay_id = f"{base_id}#{suffix}"
suffix += 1
used_frame_ids.add(overlay_id)
frames.append(
FrameOverlay(
overlay_id=overlay_id,
position=position,
wxyz=wxyz,
axes_length=axis_length,
axes_radius=axis_radius,
# Native handles report hidden in headless mode even
# though draw_marker() requested a visible marker.
visible=True,
)
)
return tuple(frames)
def _prepare_overlays(self, overlays: SceneOverlays | None) -> SceneOverlays:
reserved_frame_ids = (
{frame.overlay_id for frame in overlays.frames}
if overlays is not None
else None
)
marker_frames = self._capture_axis_marker_overlays(reserved_frame_ids)
if overlays is None:
return SceneOverlays(frames=marker_frames)
point_clouds: list[PointCloudOverlay] = []
for point_cloud in overlays.point_clouds:
point_count = point_cloud.points.shape[0]
if point_count <= self.cfg.point_cloud_max_points:
point_clouds.append(point_cloud)
continue
indices = np.linspace(
0,
point_count - 1,
num=self.cfg.point_cloud_max_points,
dtype=np.int64,
)
colors = (
point_cloud.colors[indices]
if point_cloud.colors.shape == point_cloud.points.shape
else point_cloud.colors
)
point_clouds.append(
PointCloudOverlay(
overlay_id=point_cloud.overlay_id,
points=point_cloud.points[indices],
colors=colors,
point_size=point_cloud.point_size,
visible=point_cloud.visible,
)
)
return SceneOverlays(
frames=marker_frames + overlays.frames,
trajectories=overlays.trajectories,
targets=overlays.targets,
point_clouds=tuple(point_clouds),
)
[docs]
def capture(
self,
*,
sim_step: int,
sim_time: float,
overlays: SceneOverlays | None = None,
capture_dynamic_geometry: bool = True,
) -> CaptureResult:
"""Capture one dynamic scene frame on the simulation thread.
Args:
sim_step: Current simulation step.
sim_time: Current simulation time in seconds.
overlays: Optional debug overlays.
capture_dynamic_geometry: Whether to copy soft-body and cloth vertices.
Returns:
Captured frame and producer-side copy duration.
"""
if self._scene_revision == 0:
raise RuntimeError("build_manifest() must be called before capture().")
started = perf_counter()
positions = np.empty((len(self._sources), 3), dtype=np.float32)
wxyz = np.empty((len(self._sources), 4), dtype=np.float32)
visible = self._source_visibility.copy()
if self._dynamic_source_indices.size:
positions[self._dynamic_source_indices] = self._env_offsets[
self._dynamic_env_ids
]
wxyz[self._dynamic_source_indices] = np.array(
[1.0, 0.0, 0.0, 0.0],
dtype=np.float32,
)
for batch in self._pose_batches:
if batch.asset_kind in {"rigid", "rigid_group"}:
pose = batch.asset.get_local_pose(to_matrix=False)
else:
pose = batch.asset.body_data.body_link_pose
all_positions, all_wxyz = pose_to_position_wxyz(pose)
if batch.component_indices is None:
batch_positions = all_positions[batch.env_ids]
batch_wxyz = all_wxyz[batch.env_ids]
else:
batch_positions = all_positions[
batch.env_ids,
batch.component_indices,
]
batch_wxyz = all_wxyz[
batch.env_ids,
batch.component_indices,
]
positions[batch.destination_indices] = (
batch_positions + self._env_offsets[batch.env_ids]
)
wxyz[batch.destination_indices] = batch_wxyz
dynamic_meshes: list[DynamicMeshUpdate] = []
if capture_dynamic_geometry:
dynamic_vertex_cache: dict[tuple[str, str], np.ndarray] = {}
for source in self._sources:
if not source.node.dynamic_geometry:
continue
if source.asset_key not in dynamic_vertex_cache:
if source.asset_key[0] == "soft":
vertices = source.asset.get_current_collision_vertices()
else:
vertices = source.asset.get_current_vertex_position()
dynamic_vertex_cache[source.asset_key] = _to_numpy(
vertices,
np.float32,
)
env_id = source.node.env_id
dynamic_meshes.append(
DynamicMeshUpdate(
node_id=source.node.node_id,
vertices=(
dynamic_vertex_cache[source.asset_key][env_id]
- self._env_offsets[env_id]
),
)
)
camera_pose_cache: dict[str, tuple[np.ndarray, np.ndarray]] = {}
for source in self._camera_sources:
uid = source.spec.sensor_uid
if uid in camera_pose_cache:
continue
if getattr(source.camera.cfg, "sensor_type", None) == "StereoCamera":
primary_pose, _ = source.camera.get_left_right_arena_pose()
pose = _to_numpy(primary_pose, np.float32)
else:
pose = _to_numpy(
source.camera.get_arena_pose(to_matrix=True),
np.float32,
)
if pose.shape != (self._sim.num_envs, 4, 4):
raise ValueError(
f"Camera {uid!r} arena poses must have shape "
f"({self._sim.num_envs}, 4, 4), received {pose.shape}."
)
# DexSim exposes camera poses in OpenGL convention. Viser camera
# frustums use OpenCV convention: +X right, +Y down, +Z forward.
pose[:, :3, 1:3] *= -1.0
camera_pose_cache[uid] = pose_to_position_wxyz(pose)
camera_positions = np.empty((len(self._camera_sources), 3), dtype=np.float32)
camera_wxyz = np.empty((len(self._camera_sources), 4), dtype=np.float32)
for index, source in enumerate(self._camera_sources):
all_positions, all_wxyz = camera_pose_cache[source.spec.sensor_uid]
env_id = source.spec.env_id
camera_positions[index] = all_positions[env_id] + self._env_offsets[env_id]
camera_wxyz[index] = all_wxyz[env_id]
gizmo_states: list[GizmoState] = []
for source in self._gizmo_sources:
position, quaternion = pose_to_position_wxyz(
source.gizmo.get_control_pose()
)
env_id = source.spec.env_id
gizmo_states.append(
GizmoState(
gizmo_id=source.spec.gizmo_id,
position=position[0] + self._env_offsets[env_id],
wxyz=quaternion[0],
visible=source.gizmo.is_visible(),
)
)
self._sequence += 1
frame = SceneFrame(
run_id=self.run_id,
scene_revision=self._scene_revision,
sequence=self._sequence,
sim_step=sim_step,
sim_time=sim_time,
node_ids=tuple(source.node.node_id for source in self._sources),
positions=positions,
wxyz=wxyz,
visible=visible,
camera_ids=tuple(source.spec.camera_id for source in self._camera_sources),
camera_positions=camera_positions,
camera_wxyz=camera_wxyz,
dynamic_meshes=tuple(dynamic_meshes),
gizmos=tuple(gizmo_states),
joint_controls=self._capture_joint_control_states(),
overlays=self._prepare_overlays(overlays),
)
return CaptureResult(frame=frame, capture_seconds=perf_counter() - started)
[docs]
def capture_camera_images(
self,
*,
sim_step: int,
sim_time: float,
) -> CameraImageCaptureResult:
"""Render and detach one RGB image for every exported camera instance."""
if self._scene_revision == 0:
raise RuntimeError("build_manifest() must be called before image capture.")
started = perf_counter()
images: list[CameraImage] = []
color_cache: dict[str, np.ndarray] = {}
for source in self._camera_sources:
camera = source.camera
uid = source.spec.sensor_uid
if not bool(getattr(camera.cfg, "enable_color", False)):
continue
if uid not in color_cache:
camera.update()
color = _to_numpy(camera.get_data()["color"], np.uint8)
expected = (
self._sim.num_envs,
int(camera.cfg.height),
int(camera.cfg.width),
4,
)
if color.shape != expected:
raise ValueError(
f"Camera {uid!r} color buffer must have shape {expected}, "
f"received {color.shape}."
)
color_cache[uid] = color
env_id = source.spec.env_id
images.append(
CameraImage(
camera_id=source.spec.camera_id,
image=color_cache[uid][env_id, :, :, :3],
)
)
self._image_sequence += 1
return CameraImageCaptureResult(
frame=CameraImageFrame(
run_id=self.run_id,
scene_revision=self._scene_revision,
sequence=self._image_sequence,
sim_step=sim_step,
sim_time=sim_time,
images=tuple(images),
),
capture_seconds=perf_counter() - started,
)