# ----------------------------------------------------------------------------
# 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.
# ----------------------------------------------------------------------------
"""EmbodiChain views over backend-neutral :mod:`dexsim.scene` batches.
Backend selection, handle rebinding, and topology revision tracking remain
owned by DexSim's ``Scene`` and batch classes. EmbodiChain adapts logical row
selections and its public pose convention ``(x, y, z, qx, qy, qz, qw)``;
backend-specific native compatibility work is delegated to narrow hooks.
Read-only row selection uses DexSim's public batches. Selected joint writes
scatter into reusable full-batch tensors on the target device so the control
path does not materialize DexSim selection indices on the host.
"""
from __future__ import annotations
from numbers import Integral
from typing import TYPE_CHECKING, Any, Sequence
import numpy as np
import torch
from .base import ArticulationViewBase, RigidBodyViewBase
if TYPE_CHECKING:
from dexsim.scene import (
ArticulationBatch,
RigidBodyBatch,
Scene,
SpawnedArticulation,
SpawnedRigidBody,
)
__all__ = ["SceneArticulationView", "SceneRigidBodyView"]
_NEWTON_ROOT_POSE_ATOL = 1.0e-6
def _checked_batch_call(
batch: Any,
method_name: str,
/,
*args: Any,
**kwargs: Any,
) -> Any:
"""Call one Scene batch operation and reject native failure statuses."""
status = getattr(batch, method_name)(*args, **kwargs)
if isinstance(status, Integral) and status < 0:
raise RuntimeError(
f"DexSim Scene batch operation {method_name!r} failed with "
f"status {status}."
)
return status
def _selection_values(
selection: Sequence[int] | torch.Tensor | None,
valid_values: torch.Tensor,
*,
label: str,
) -> torch.Tensor:
if selection is None:
return valid_values
result = torch.as_tensor(
selection,
dtype=torch.long,
device=valid_values.device,
).reshape(-1)
try:
# ``index_select`` validates bounds on the tensor's device. Unlike a
# Python boolean reduction, this does not synchronize CUDA selections
# back to the host on every control write.
return torch.index_select(valid_values, 0, result)
except IndexError:
raise IndexError(
f"{label} selection is outside [0, {len(valid_values)})."
) from None
def _batch_pose(data: torch.Tensor) -> torch.Tensor:
"""Convert rigid-body ``xyz+xyzw`` poses to batch ``xyzw+xyz``."""
result = torch.empty_like(data, dtype=torch.float32)
result[..., 0:4] = data[..., 3:7]
result[..., 4:7] = data[..., 0:3]
return result
def _embodichain_pose(data: torch.Tensor) -> torch.Tensor:
"""Convert batch ``xyzw+xyz`` poses to rigid-body ``xyz+xyzw``."""
result = torch.empty_like(data, dtype=torch.float32)
result[..., 0:3] = data[..., 4:7]
result[..., 3:7] = data[..., 0:4]
return result
class _SceneVisualObjectView:
"""Adapt pure-visual World batches from world to arena-local matrices.
Parent transforms and native IDs are captured once per topology revision.
Steady-state reads and writes each use one synchronous native tensor call,
including selections on CUDA. Outputs belong to the caller.
"""
def __init__(
self, scene: Scene, entities: Sequence[SpawnedRigidBody], device: torch.device
) -> None:
self.scene = scene
self.device = device
self.refresh(entities)
def refresh(self, entities: Sequence[SpawnedRigidBody]) -> None:
self.rows = torch.arange(len(entities), dtype=torch.long, device=self.device)
self.user_ids = torch.tensor(
[entity.get_user_id() for entity in entities],
dtype=torch.int32,
device=self.device,
)
local = torch.as_tensor(
np.stack([entity.get_local_pose() for entity in entities]),
dtype=torch.float32,
device=self.device,
)
world = torch.as_tensor(
np.stack([entity.get_world_pose() for entity in entities]),
dtype=torch.float32,
device=self.device,
)
self._parent_world = world @ torch.linalg.inv(local)
self._parent_inverse = torch.linalg.inv(self._parent_world)
def fetch_pose(self, rows: torch.Tensor) -> torch.Tensor:
transforms = torch.empty(
(len(rows), 4, 4), dtype=torch.float32, device=self.device
)
if len(rows):
_checked_batch_call(
self.scene.world,
"fetch_object_transform",
transforms,
self.user_ids.index_select(0, rows),
)
transforms = self._parent_inverse.index_select(0, rows) @ transforms
return transforms
def apply_pose(self, poses: torch.Tensor, rows: torch.Tensor) -> None:
if len(rows):
transforms = self._parent_world.index_select(0, rows) @ poses
_checked_batch_call(
self.scene.world,
"apply_object_transform",
transforms.contiguous(),
self.user_ids.index_select(0, rows),
)
class _SceneBatchSelectionAdapter:
"""Shared row-selection support for fixed-size Scene batches."""
def __init__(self, batch: Any, device: torch.device, row_count: int) -> None:
self.batch = batch
self.device = device
self._row_count = row_count
self._row_indices = torch.arange(
row_count,
dtype=torch.long,
device=device,
)
def _select_rows(
self,
selection: Sequence[int] | torch.Tensor | None,
) -> torch.Tensor:
return _selection_values(
selection,
self._row_indices,
label="Batch row",
)
def _fetch_rows(
self,
method_name: str,
out: torch.Tensor,
selection: Sequence[int] | torch.Tensor | None,
tail_shape: tuple[int, ...],
) -> torch.Tensor:
if selection is None:
expected_shape = (self._row_count, *tail_shape)
if tuple(out.shape) != expected_shape:
raise ValueError(
f"Expected batch output shape {expected_shape}, got "
f"{tuple(out.shape)}."
)
_checked_batch_call(self.batch, method_name, out)
return out
rows = self._select_rows(selection)
expected_shape = (len(rows), *tail_shape)
if tuple(out.shape) != expected_shape:
raise ValueError(
f"Expected selected output shape {expected_shape}, got "
f"{tuple(out.shape)}."
)
if len(rows):
_checked_batch_call(self.batch.select(rows), method_name, out)
return out
def _apply_rows(
self,
method_name: str,
values: torch.Tensor,
selection: Sequence[int] | torch.Tensor,
tail_shape: tuple[int, ...],
) -> None:
rows = self._select_rows(selection)
expected_shape = (len(rows), *tail_shape)
if tuple(values.shape) != expected_shape:
raise ValueError(
f"Expected selected data shape {expected_shape}, got "
f"{tuple(values.shape)}."
)
if len(rows):
_checked_batch_call(self.batch.select(rows), method_name, values)
[docs]
class SceneRigidBodyView(_SceneBatchSelectionAdapter, RigidBodyViewBase):
"""Backend-neutral rigid-body view backed by ``RigidBodyBatch``."""
[docs]
@classmethod
def from_entities(
cls,
scene: Scene,
entities: Sequence[SpawnedRigidBody],
device: torch.device,
) -> SceneRigidBodyView:
"""Create a view from rigid-body handles owned by ``scene``."""
return cls(
scene,
scene.create_rigid_body_batch(list(entities)),
device,
)
[docs]
def __init__(
self,
scene: Scene,
batch: RigidBodyBatch,
device: torch.device,
) -> None:
super().__init__(batch, device, len(batch))
self.scene = scene
self._body_ids_tensor = torch.arange(
len(batch), dtype=torch.int32, device=device
)
self._newton_state_sync: tuple[int, Any, Any] | None = None
@property
def is_ready(self) -> bool:
return True
@property
def is_newton_backend(self) -> bool:
return self.scene.backend == "newton"
@property
def body_ids(self) -> list[int]:
return list(range(self._row_count))
@property
def body_ids_tensor(self) -> torch.Tensor:
return self._body_ids_tensor
[docs]
def select_body_ids(self, indices: Sequence[int] | torch.Tensor) -> torch.Tensor:
return self._body_ids_tensor[indices]
[docs]
def fetch_pose(
self, data: torch.Tensor, body_ids: torch.Tensor | None = None
) -> None:
batch_pose = torch.empty(
(len(data), 7), dtype=torch.float32, device=self.device
)
self._fetch_rows("fetch_pose", batch_pose, body_ids, (7,))
data.copy_(_embodichain_pose(batch_pose).to(data.device, data.dtype))
[docs]
def apply_pose(self, pose: torch.Tensor, body_ids: torch.Tensor) -> None:
self._apply_rows(
"apply_pose",
_batch_pose(pose.to(self.device, torch.float32)),
body_ids,
(7,),
)
if self.is_newton_backend and len(body_ids):
self._synchronize_newton_standalone_state()
def _synchronize_newton_standalone_state(self) -> None:
"""Keep Newton standalone-body FREE joints coherent after state writes.
DexSim 0.4.3's device ``RigidBodyBatch`` state writes update maximal
``body_q`` or ``body_qd`` state, while MuJoCo-Warp advances standalone
rigid bodies from their reduced FREE-joint state. Cache one selection
for this stable batch and project both state buffers after each write.
"""
from .newton import _synchronize_standalone_rigid_body_state
self._newton_state_sync = _synchronize_standalone_rigid_body_state(
self.scene,
self.batch,
self._newton_state_sync,
)
[docs]
def fetch_com_local_pose(
self, data: torch.Tensor, body_ids: torch.Tensor | None = None
) -> None:
batch_pose = torch.empty(
(len(data), 7), dtype=torch.float32, device=self.device
)
self._fetch_rows("fetch_com_local_pose", batch_pose, body_ids, (7,))
data.copy_(_embodichain_pose(batch_pose).to(data.device, data.dtype))
[docs]
def apply_com_local_pose(self, data: torch.Tensor, body_ids: torch.Tensor) -> None:
self._apply_rows(
"apply_com_local_pose",
_batch_pose(data.to(self.device, torch.float32)),
body_ids,
(7,),
)
[docs]
def fetch_linear_velocity(
self, data: torch.Tensor, body_ids: torch.Tensor | None = None
) -> None:
self._fetch_rows("fetch_linear_velocity", data, body_ids, (3,))
[docs]
def fetch_angular_velocity(
self, data: torch.Tensor, body_ids: torch.Tensor | None = None
) -> None:
self._fetch_rows("fetch_angular_velocity", data, body_ids, (3,))
[docs]
def apply_linear_velocity(self, data: torch.Tensor, body_ids: torch.Tensor) -> None:
self._apply_rows(
"apply_linear_velocity",
data,
body_ids,
(3,),
)
if self.is_newton_backend and len(body_ids):
self._synchronize_newton_standalone_state()
[docs]
def apply_angular_velocity(
self, data: torch.Tensor, body_ids: torch.Tensor
) -> None:
self._apply_rows(
"apply_angular_velocity",
data,
body_ids,
(3,),
)
if self.is_newton_backend and len(body_ids):
self._synchronize_newton_standalone_state()
[docs]
def fetch_linear_acceleration(
self, data: torch.Tensor, body_ids: torch.Tensor | None = None
) -> None:
self._fetch_rows("fetch_linear_acceleration", data, body_ids, (3,))
[docs]
def fetch_angular_acceleration(
self, data: torch.Tensor, body_ids: torch.Tensor | None = None
) -> None:
self._fetch_rows("fetch_angular_acceleration", data, body_ids, (3,))
[docs]
def apply_force(self, data: torch.Tensor, body_ids: torch.Tensor) -> None:
self._apply_rows("apply_force", data, body_ids, (3,))
[docs]
def apply_torque(self, data: torch.Tensor, body_ids: torch.Tensor) -> None:
self._apply_rows("apply_torque", data, body_ids, (3,))
[docs]
def fetch_mass(
self, data: torch.Tensor, body_ids: torch.Tensor | None = None
) -> None:
self._fetch_rows("fetch_mass", data, body_ids, (1,))
[docs]
def apply_mass(self, data: torch.Tensor, body_ids: torch.Tensor) -> None:
self._apply_rows("apply_mass", data, body_ids, (1,))
[docs]
def fetch_inertia_diagonal(
self, data: torch.Tensor, body_ids: torch.Tensor | None = None
) -> None:
self._fetch_rows("fetch_inertia_diagonal", data, body_ids, (3,))
[docs]
def apply_inertia_diagonal(
self, data: torch.Tensor, body_ids: torch.Tensor
) -> None:
self._apply_rows(
"apply_inertia_diagonal",
data,
body_ids,
(3,),
)
[docs]
def fetch_friction(
self, data: torch.Tensor, body_ids: torch.Tensor | None = None
) -> None:
self._fetch_rows("fetch_friction", data, body_ids, (1,))
[docs]
def apply_friction(self, data: torch.Tensor, body_ids: torch.Tensor) -> None:
self._apply_rows("apply_friction", data, body_ids, (1,))
[docs]
def fetch_restitution(
self, data: torch.Tensor, body_ids: torch.Tensor | None = None
) -> None:
self._fetch_rows("fetch_restitution", data, body_ids, (1,))
[docs]
def apply_restitution(self, data: torch.Tensor, body_ids: torch.Tensor) -> None:
self._apply_rows("apply_restitution", data, body_ids, (1,))
[docs]
def fetch_contact_offset(
self, data: torch.Tensor, body_ids: torch.Tensor | None = None
) -> None:
self._fetch_rows("fetch_contact_offset", data, body_ids, (1,))
[docs]
def apply_contact_offset(self, data: torch.Tensor, body_ids: torch.Tensor) -> None:
self._apply_rows("apply_contact_offset", data, body_ids, (1,))
[docs]
def fetch_damping(
self, data: torch.Tensor, body_ids: torch.Tensor | None = None
) -> None:
self._fetch_rows("fetch_damping", data, body_ids, (2,))
[docs]
def apply_damping(self, data: torch.Tensor, body_ids: torch.Tensor) -> None:
self._apply_rows("apply_damping", data, body_ids, (2,))
[docs]
def fetch_collision_filter(
self,
data: torch.Tensor,
body_ids: torch.Tensor | None = None,
) -> None:
self._fetch_rows("fetch_collision_filter", data, body_ids, (4,))
[docs]
def apply_collision_filter(
self,
data: torch.Tensor,
body_ids: torch.Tensor,
) -> None:
self._apply_rows("apply_collision_filter", data, body_ids, (4,))
[docs]
class SceneArticulationView(_SceneBatchSelectionAdapter, ArticulationViewBase):
"""Backend-neutral articulation state view backed by ``ArticulationBatch``.
Joint selections currently require one scalar DOF per selected joint. The
public DexSim layout already describes multi-DOF joints; supporting them
without ambiguity requires a DOF-selection API in DexSim and is therefore
kept as an explicit boundary rather than guessed here.
"""
[docs]
@classmethod
def from_entities(
cls,
scene: Scene,
entities: Sequence[SpawnedArticulation],
device: torch.device,
) -> SceneArticulationView:
"""Create a view from articulation handles owned by ``scene``."""
return cls(
scene,
scene.create_articulation_batch(list(entities)),
device,
)
[docs]
def __init__(
self,
scene: Scene,
batch: ArticulationBatch,
device: torch.device,
) -> None:
super().__init__(batch, device, len(batch))
self.scene = scene
self._validate_homogeneous_layout()
layouts = tuple(self.batch.joint_layouts_per_articulation)
self._joint_dof_columns_host = tuple(
[layout.dof_start for layout in layouts[0]]
if layouts
else list(range(self.batch.dof_width))
)
self._joint_dof_columns = torch.as_tensor(
self._joint_dof_columns_host,
dtype=torch.long,
device=device,
)
self._articulation_ids = torch.arange(
len(batch), dtype=torch.int32, device=device
)
self._link_com_scratch: torch.Tensor | None = None
self._root_velocity_snapshot: torch.Tensor | None = None
self._root_pose_scratch: torch.Tensor | None = None
self._joint_apply_scratch: dict[str, torch.Tensor] = {}
def _validate_homogeneous_layout(self) -> None:
"""Require the uniform topology promised by one EC Articulation."""
dof_counts = tuple(self.batch.dof_counts)
link_counts = tuple(self.batch.link_counts)
joint_names = tuple(self.batch.joint_names_per_articulation)
link_names = tuple(self.batch.link_names_per_articulation)
if dof_counts and len(set(dof_counts)) != 1:
raise ValueError(
"One EmbodiChain Articulation cannot bind heterogeneous Scene "
f"DOF counts: {dof_counts}."
)
if link_counts and len(set(link_counts)) != 1:
raise ValueError(
"One EmbodiChain Articulation cannot bind heterogeneous Scene "
f"link counts: {link_counts}."
)
if joint_names and any(names != joint_names[0] for names in joint_names[1:]):
raise ValueError(
"One EmbodiChain Articulation requires identical active-joint "
"ordering in every Scene row."
)
if link_names and any(names != link_names[0] for names in link_names[1:]):
raise ValueError(
"One EmbodiChain Articulation requires identical link ordering "
"in every Scene row."
)
layouts = tuple(self.batch.joint_layouts_per_articulation)
if layouts and any(layout.dof_count != 1 for layout in layouts[0]):
raise NotImplementedError(
"EmbodiChain's Articulation API currently indexes joints and "
"scalar DOFs interchangeably. Scene multi-DOF joints require "
"an explicit DOF-selection API before they can be bound safely."
)
@property
def dof(self) -> int:
"""Scalar DOF width shared by every articulation row."""
return self.batch.dof_width
@property
def num_links(self) -> int:
"""Link count shared by every articulation row."""
return self.batch.link_width
@property
def joint_names(self) -> list[str]:
"""Active joints in public flattened-DOF order."""
rows = self.batch.joint_names_per_articulation
return [] if not rows else list(rows[0])
@property
def link_names(self) -> list[str]:
"""Links in public link-buffer order."""
rows = self.batch.link_names_per_articulation
return [] if not rows else list(rows[0])
@property
def is_ready(self) -> bool:
return True
@property
def is_newton_backend(self) -> bool:
return self.scene.backend == "newton"
@property
def articulation_ids_tensor(self) -> torch.Tensor:
return self._articulation_ids
[docs]
def select_articulation_ids(
self, env_ids: Sequence[int] | torch.Tensor
) -> torch.Tensor:
return self._articulation_ids[env_ids]
[docs]
def fetch_link_physical_properties(
self,
mass: torch.Tensor,
inertia: torch.Tensor,
com_pose: torch.Tensor,
) -> None:
"""Read live batch properties into caller-owned output buffers.
Args:
mass: Output masses, shape ``(N, num_links)``.
inertia: Output principal moments, shape ``(N, num_links, 3)``.
com_pose: Output local poses, shape ``(N, num_links, 7)``,
in EmbodiChain's ``xyz + xyzw`` convention.
"""
scratch = self._link_com_scratch
if scratch is None or scratch.shape != com_pose.shape:
scratch = torch.empty_like(
com_pose, dtype=torch.float32, device=self.device
)
self._link_com_scratch = scratch
_checked_batch_call(self.batch, "fetch_link_mass", mass)
_checked_batch_call(self.batch, "fetch_link_inertia_diagonal", inertia)
_checked_batch_call(self.batch, "fetch_link_com_local_pose", scratch)
com_pose[..., :3].copy_(scratch[..., 4:])
com_pose[..., 3:].copy_(scratch[..., :4])
[docs]
def fetch_joint_properties(
self,
*,
position_limits: torch.Tensor | None = None,
velocity_limit: torch.Tensor | None = None,
effort_limit: torch.Tensor | None = None,
stiffness: torch.Tensor | None = None,
damping: torch.Tensor | None = None,
friction: torch.Tensor | None = None,
armature: torch.Tensor | None = None,
) -> None:
"""Read current parameters without tracking model gradients.
See :meth:`ArticulationViewBase.fetch_joint_properties` for buffer
shapes, units and lifetime.
"""
# Configuration snapshots must not inherit gradients from Newton's
# differentiable model parameters.
with torch.no_grad():
_checked_batch_call(
self.batch,
"fetch_joint_properties",
position_limits=position_limits,
velocity_limit=velocity_limit,
effort_limit=effort_limit,
stiffness=stiffness,
damping=damping,
friction=friction,
armature=armature,
)
[docs]
def fetch_state(
self,
qpos: torch.Tensor,
qvel: torch.Tensor,
root_pose: torch.Tensor,
root_lin_vel: torch.Tensor,
root_ang_vel: torch.Tensor,
) -> None:
"""Read joint and root state through DexSim's existing batch operation.
Buffer shapes and frames follow :meth:`ArticulationViewBase.fetch_state`.
Each call reads current simulation data, including writes since the
previous call; the reusable scratch buffer stores no cached state.
"""
scratch = self._root_pose_scratch
if scratch is None or scratch.shape != root_pose.shape:
scratch = torch.empty_like(
root_pose, dtype=torch.float32, device=self.device
)
self._root_pose_scratch = scratch
_checked_batch_call(
self.batch, "fetch_state", qpos, qvel, scratch, root_lin_vel, root_ang_vel
)
root_pose[..., :3].copy_(scratch[..., 4:])
root_pose[..., 3:].copy_(scratch[..., :4])
[docs]
def fetch_root_pose(self, data: torch.Tensor) -> torch.Tensor:
batch_pose = torch.empty_like(data, dtype=torch.float32, device=self.device)
_checked_batch_call(self.batch, "fetch_root_pose", batch_pose)
data.copy_(_embodichain_pose(batch_pose).to(data.device, data.dtype))
return data
[docs]
def fetch_root_linear_velocity(self, data: torch.Tensor) -> torch.Tensor:
_checked_batch_call(self.batch, "fetch_root_linear_velocity", data)
return data
[docs]
def fetch_root_angular_velocity(self, data: torch.Tensor) -> torch.Tensor:
_checked_batch_call(self.batch, "fetch_root_angular_velocity", data)
return data
[docs]
def fetch_qpos(self, data: torch.Tensor) -> torch.Tensor:
_checked_batch_call(self.batch, "fetch_joint_position", data)
return data
[docs]
def fetch_target_qpos(self, data: torch.Tensor) -> torch.Tensor:
_checked_batch_call(self.batch, "fetch_joint_target_position", data)
return data
[docs]
def fetch_qvel(self, data: torch.Tensor) -> torch.Tensor:
_checked_batch_call(self.batch, "fetch_joint_velocity", data)
return data
[docs]
def fetch_target_qvel(self, data: torch.Tensor) -> torch.Tensor:
_checked_batch_call(self.batch, "fetch_joint_target_velocity", data)
return data
[docs]
def fetch_qacc(self, data: torch.Tensor) -> torch.Tensor:
_checked_batch_call(self.batch, "fetch_joint_acceleration", data)
return data
[docs]
def fetch_qf(self, data: torch.Tensor) -> torch.Tensor:
_checked_batch_call(self.batch, "fetch_joint_force", data)
return data
[docs]
def fetch_link_pose(self, data: torch.Tensor) -> torch.Tensor:
batch_pose = torch.empty_like(data, dtype=torch.float32, device=self.device)
_checked_batch_call(self.batch, "fetch_link_pose", batch_pose)
data.copy_(_embodichain_pose(batch_pose).to(data.device, data.dtype))
return data
[docs]
def fetch_link_velocity(
self,
data: torch.Tensor,
linear_data: torch.Tensor,
angular_data: torch.Tensor,
) -> torch.Tensor:
_checked_batch_call(self.batch, "fetch_link_linear_velocity", linear_data)
_checked_batch_call(self.batch, "fetch_link_angular_velocity", angular_data)
data[..., 0:3] = linear_data
data[..., 3:6] = angular_data
return data
[docs]
def apply_root_pose(
self, pose: torch.Tensor, env_ids: Sequence[int] | torch.Tensor
) -> None:
rows = self._select_rows(env_ids)
batch_pose = _batch_pose(pose.to(self.device, torch.float32))
expected_shape = (len(rows), 7)
if tuple(batch_pose.shape) != expected_shape:
raise ValueError(
f"Expected selected data shape {expected_shape}, got "
f"{tuple(batch_pose.shape)}."
)
if self.is_newton_backend and len(rows):
current_pose = torch.empty_like(batch_pose)
self._fetch_rows(
"fetch_root_pose",
current_pose,
rows,
(7,),
)
translation_matches = torch.all(
torch.abs(current_pose[:, 4:7] - batch_pose[:, 4:7])
<= _NEWTON_ROOT_POSE_ATOL,
dim=1,
)
quaternion_delta = torch.minimum(
torch.amax(torch.abs(current_pose[:, 0:4] - batch_pose[:, 0:4]), dim=1),
torch.amax(torch.abs(current_pose[:, 0:4] + batch_pose[:, 0:4]), dim=1),
)
changed = ~(
translation_matches & (quaternion_delta <= _NEWTON_ROOT_POSE_ATOL)
)
rows = rows[changed]
batch_pose = batch_pose[changed]
if len(rows):
_checked_batch_call(
self.batch, "apply_state", rows=rows, root_pose=batch_pose
)
if not self.is_newton_backend:
# Default root writes do not propagate descendant poses on CUDA.
self.compute_kinematics(rows)
[docs]
def apply_root_velocity(
self, velocity: torch.Tensor, env_ids: Sequence[int] | torch.Tensor
) -> None:
"""Write selected root velocities through the existing Scene batch.
Args:
velocity: World-frame linear and angular velocities, shape ``(N, 6)``.
env_ids: Environment rows in the same order as ``velocity``.
Raises:
RuntimeError: A batch write fails. Both velocity components are
restored first on the default backend; a failed rollback is
reported with the original write error as its cause. Newton
uses the batch state writer, which validates both components
before writing and does not roll back execution failures.
"""
if getattr(self.scene, "backend", None) == "newton":
# The batch state writer validates both components before mutation
# and propagates the final state once. No per-row batch or separate
# snapshot/linear/angular writes are needed for this backend.
self.apply_state(env_ids=env_ids, root_velocity=velocity)
return
rows = self._select_rows(env_ids)
expected_shape = (len(rows), 6)
if tuple(velocity.shape) != expected_shape:
raise ValueError(
f"Expected selected data shape {expected_shape}, got "
f"{tuple(velocity.shape)}."
)
if len(rows):
batch = self.batch.select(rows)
if (
self._root_velocity_snapshot is None
or self._root_velocity_snapshot.shape[1] < len(rows)
):
self._root_velocity_snapshot = torch.empty(
(2, len(rows), 3), dtype=torch.float32, device=self.device
)
previous = self._root_velocity_snapshot[:, : len(rows)]
# Snapshot both components before either independent native write.
_checked_batch_call(batch, "fetch_root_linear_velocity", previous[0])
_checked_batch_call(batch, "fetch_root_angular_velocity", previous[1])
try:
_checked_batch_call(
batch, "apply_root_linear_velocity", velocity[:, :3]
)
_checked_batch_call(
batch, "apply_root_angular_velocity", velocity[:, 3:]
)
except Exception as write_error:
rollback_errors = []
for component, saved in zip(("linear", "angular"), previous):
method = f"apply_root_{component}_velocity"
try:
_checked_batch_call(batch, method, saved)
except Exception as rollback_error:
rollback_errors.append(f"{method}: {rollback_error}")
if rollback_errors:
raise RuntimeError(
"Root velocity write failed and rollback could not restore "
"all selected velocities: " + "; ".join(rollback_errors)
) from write_error
raise
def _joint_columns(
self, joint_ids: Sequence[int] | torch.Tensor | None
) -> torch.Tensor | tuple[int, ...] | None:
# Static joint selections originate on the host. Keep them there for
# DexSim's topology validation instead of uploading then downloading
# the same indices on every control command.
if not isinstance(joint_ids, torch.Tensor):
columns = self._joint_dof_columns_host
if joint_ids is not None:
ids = tuple(int(index) for index in joint_ids)
if any(index < 0 or index >= len(columns) for index in ids):
raise IndexError(f"Joint selection is outside [0, {len(columns)}).")
columns = tuple(columns[index] for index in ids)
return None if columns == tuple(range(self.dof)) else columns
return _selection_values(
joint_ids,
self._joint_dof_columns,
label="Joint",
)
[docs]
def apply_state(
self,
env_ids: Sequence[int] | torch.Tensor | None = None,
joint_ids: Sequence[int] | torch.Tensor | None = None,
*,
root_pose: torch.Tensor | None = None,
qpos: torch.Tensor | None = None,
target_qpos: torch.Tensor | None = None,
qvel: torch.Tensor | None = None,
target_qvel: torch.Tensor | None = None,
qf: torch.Tensor | None = None,
root_velocity: torch.Tensor | None = None,
clear_dynamics: bool = False,
) -> None:
"""Map EC fields once and submit one selected Scene batch write.
Shapes and selection semantics follow
:meth:`ArticulationViewBase.apply_state`.
"""
if not isinstance(env_ids, torch.Tensor) and env_ids is not None:
if tuple(env_ids) == tuple(range(self._row_count)):
env_ids = None
rows = None if env_ids is None else self._select_rows(env_ids)
count = self._row_count if rows is None else len(rows)
if count == 0:
return
fields = {
name: value
for name, value in (
("joint_position", qpos),
("joint_target_position", target_qpos),
("joint_velocity", qvel),
("joint_target_velocity", target_qvel),
("joint_force", qf),
)
if value is not None
}
columns = self._joint_columns(joint_ids) if fields else None
if (
fields
and root_pose is None
and root_velocity is None
and not clear_dynamics
and (rows is not None or isinstance(columns, torch.Tensor))
):
# Selected control commands use the existing full-batch scatter
# path: DexSim otherwise downloads tensor row/DOF selections.
selected_rows = self._row_indices if rows is None else rows
selected_columns = (
self._joint_dof_columns
if columns is None
else torch.as_tensor(columns, dtype=torch.long, device=self.device)
)
expected = (len(selected_rows), len(selected_columns))
for name, value in fields.items():
if tuple(value.shape) != expected:
raise ValueError(
f"Expected {name} shape {expected}, got {tuple(value.shape)}."
)
if not len(selected_columns):
return
for name, value in tuple(fields.items()):
scratch = self._joint_apply_scratch.get(name)
if scratch is None:
scratch = torch.empty(
(self._row_count, self.dof),
dtype=torch.float32,
device=self.device,
)
self._joint_apply_scratch[name] = scratch
_checked_batch_call(self.batch, f"fetch_{name}", scratch)
scratch[selected_rows[:, None], selected_columns] = value.to(
self.device, torch.float32
)
fields[name] = scratch
rows = columns = None
if root_pose is not None:
if tuple(root_pose.shape) != (count, 7):
raise ValueError(f"Expected root_pose shape {(count, 7)}.")
fields["root_pose"] = _batch_pose(root_pose.to(self.device, torch.float32))
if root_velocity is not None:
if tuple(root_velocity.shape) != (count, 6):
raise ValueError(f"Expected root_velocity shape {(count, 6)}.")
fields["root_linear_velocity"] = root_velocity[:, :3]
fields["root_angular_velocity"] = root_velocity[:, 3:]
propagate_root = (
root_pose is not None
and qpos is None
and not clear_dynamics
and not self.is_newton_backend
)
batch = self.batch
if propagate_root and rows is not None:
batch = batch.select(rows)
rows = None
_checked_batch_call(
batch,
"apply_state",
rows=rows,
dof_ids=columns,
clear_dynamics=clear_dynamics,
**fields,
)
if propagate_root:
# Default's native state writer propagates qpos/clear operations,
# but root-only writes need an explicit descendant update.
_checked_batch_call(batch, "compute_kinematics")
[docs]
def apply_qpos(
self,
qpos: torch.Tensor,
env_ids: Sequence[int] | torch.Tensor | None,
joint_ids: Sequence[int] | torch.Tensor | None,
*,
target: bool,
) -> None:
self.apply_state(
env_ids, joint_ids, **{"target_qpos" if target else "qpos": qpos}
)
[docs]
def apply_qvel(
self,
qvel: torch.Tensor,
env_ids: Sequence[int] | torch.Tensor | None,
joint_ids: Sequence[int] | torch.Tensor | None,
*,
target: bool,
) -> None:
self.apply_state(
env_ids, joint_ids, **{"target_qvel" if target else "qvel": qvel}
)
[docs]
def apply_qf(
self,
qf: torch.Tensor,
env_ids: Sequence[int] | torch.Tensor | None,
joint_ids: Sequence[int] | torch.Tensor | None,
) -> None:
self.apply_state(env_ids, joint_ids, qf=qf)
[docs]
def clear_dynamics(self, env_ids: Sequence[int] | torch.Tensor) -> None:
self.apply_state(env_ids, clear_dynamics=True)
[docs]
def compute_kinematics(self, env_ids: Sequence[int] | torch.Tensor) -> None:
rows = self._select_rows(env_ids)
if not len(rows):
return
_checked_batch_call(self.batch.select(rows), "compute_kinematics")