# ----------------------------------------------------------------------------
# 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 torch
import numpy as np
from dataclasses import dataclass
from scipy.spatial.transform import Rotation, Slerp
from enum import Enum
from typing import Union, List
from embodichain.utils import logger
__all__ = [
"TrajectorySampleMethod",
"MovePart",
"MoveType",
"PlanState",
"PlanResult",
"normalize_success_mask",
"calculate_point_allocations",
"interpolate_xpos",
"interpolate_xpos_batched",
]
def normalize_success_mask(
success: bool | torch.Tensor,
*,
num_envs: int,
device: torch.device | str,
name: str,
) -> torch.Tensor:
"""Normalize a scalar or batched success value to ``(num_envs,)``.
Args:
success: Scalar success or a boolean/binary-integer tensor.
num_envs: Required batch size.
device: Device of the resulting tensor.
name: Human-readable value name used in validation errors.
Returns:
Independently owned boolean success mask.
Raises:
TypeError: If ``success`` is neither boolean nor binary integer data.
ValueError: If a tensor does not match the required batch shape or a
CUDA device is requested while CUDA is unavailable.
"""
resolved_device = torch.device(device)
if resolved_device.type == "cuda":
if not torch.cuda.is_available():
raise ValueError(
"CUDA device requested for success-mask normalization, but "
"torch.cuda.is_available() is False."
)
if resolved_device.index is None:
resolved_device = torch.device(f"cuda:{torch.cuda.current_device()}")
if isinstance(success, bool):
return torch.full(
(num_envs,), success, dtype=torch.bool, device=resolved_device
)
if not isinstance(success, torch.Tensor):
raise TypeError(
f"{name} must be a bool or torch.Tensor, got {type(success).__name__}."
)
success = success.to(resolved_device)
if success.dtype != torch.bool:
integer_dtypes = {
torch.uint8,
torch.int8,
torch.int16,
torch.int32,
torch.int64,
}
if success.dtype not in integer_dtypes or not torch.all(
(success == 0) | (success == 1)
):
raise TypeError(
f"{name} must be boolean or a binary integer tensor, "
f"got dtype {success.dtype}."
)
success = success.to(dtype=torch.bool)
if success.dim() == 0 or success.shape == (1,):
success = success.reshape(1).expand(num_envs)
if success.shape != (num_envs,):
raise ValueError(
f"{name} must have shape ({num_envs},), got {tuple(success.shape)}."
)
return success.clone()
[docs]
class TrajectorySampleMethod(Enum):
r"""Enumeration for different trajectory sampling methods.
This enum defines various methods for sampling trajectories,
providing meaningful names for different sampling strategies.
"""
TIME = "time"
"""Sample based on time intervals."""
QUANTITY = "quantity"
"""Sample based on a specified number of points."""
DISTANCE = "distance"
"""Sample based on distance intervals."""
[docs]
@classmethod
def from_str(
cls, value: Union[str, "TrajectorySampleMethod"]
) -> "TrajectorySampleMethod":
if isinstance(value, cls):
return value
try:
return cls[value.upper()]
except KeyError:
valid_values = [e.name for e in cls]
logger.log_error(
f"Invalid version '{value}'. Valid values are: {valid_values}",
ValueError,
)
def __str__(self):
"""Override string representation for better readability."""
return self.value.capitalize()
[docs]
class MovePart(Enum):
r"""Enumeration for different robot parts to move.
Defines robot part selection for motion planning.
Attributes:
LEFT (int): left arm or end-effector.
RIGHT (int): right arm or end-effector.
BOTH (int): both arms or end-effectors.
TORSO (int): torso for humanoid robot.
ALL (int): all joints of the robot (joint control only).
"""
LEFT = 0 # left arm|eef
RIGHT = 1 # right arm|eef
BOTH = 2 # left arm|eef and right arm|eef
TORSO = 3 # torso for humanoid robot
ALL = 4 # all joints of the robot. Only for joint control.
[docs]
class MoveType(Enum):
r"""Enumeration for different types of movements.
Defines movement types for robot planning.
Attributes:
TOOL (int): Tool open or close.
EEF_MOVE (int): Move end-effector to target pose (IK + trajectory).
JOINT_MOVE (int): Move joints to target angles (trajectory planning).
SYNC (int): Synchronized left/right arm movement (dual-arm robots).
PAUSE (int): Pause for specified duration (see PlanState.pause_seconds).
"""
TOOL = 0 # Tool open or close
EEF_MOVE = 1 # Move the end-effector to a target pose (xpos) using IK and trajectory planning
JOINT_MOVE = (
2 # Directly move joints to target angles (qpos) using trajectory planning
)
SYNC = 3 # Synchronized left and right arm movement (for dual-arm robots)
PAUSE = 4 # Pause for a specified duration (use pause_seconds in PlanState)
[docs]
@dataclass
class PlanResult:
r"""Data class representing the result of a motion plan (env-batched).
A result that contains joint positions must also contain per-sample ``dt``.
Per-environment :attr:`duration` is derived from those intervals. Failed
plans may omit all trajectory fields by leaving ``positions`` as ``None``.
"""
success: bool | torch.Tensor = False
"""Per-env success, shape ``(B,)`` bool tensor (or scalar bool)."""
xpos_list: torch.Tensor | None = None
"""End-effector poses, shape ``(B, N, 4, 4)``."""
positions: torch.Tensor | None = None
"""Joint positions, shape ``(B, N, DOF)``."""
velocities: torch.Tensor | None = None
"""Joint velocities, shape ``(B, N, DOF)``."""
accelerations: torch.Tensor | None = None
"""Joint accelerations, shape ``(B, N, DOF)``."""
dt: torch.Tensor | None = None
"""Per-env time deltas, shape ``(B, N)``."""
def __post_init__(self) -> None:
"""Validate the explicit trajectory-timing contract."""
if self.positions is None:
if self.dt is not None:
raise ValueError("PlanResult timing requires positions.")
return
if not isinstance(self.positions, torch.Tensor) or self.positions.dim() != 3:
raise ValueError("PlanResult.positions must have shape (B, N, DOF).")
batch_size, waypoint_count, _ = self.positions.shape
if not isinstance(self.dt, torch.Tensor):
raise ValueError(
"PlanResult with positions requires explicit dt with shape (B, N)."
)
if self.dt.shape != (batch_size, waypoint_count):
raise ValueError(
"PlanResult.dt must match positions batch and waypoint dimensions."
)
if self.dt.device != self.positions.device:
raise ValueError("PlanResult.dt and positions must share a device.")
if not torch.isfinite(self.dt).all() or (self.dt < 0).any():
raise ValueError("PlanResult.dt must contain finite non-negative values.")
@property
def duration(self) -> torch.Tensor | None:
"""Return per-environment duration derived from :attr:`dt`."""
return None if self.dt is None else self.dt.sum(dim=1)
[docs]
def is_all_success(self) -> bool:
"""Return True only when every env succeeded."""
if isinstance(self.success, torch.Tensor):
return bool(torch.all(self.success).item())
return bool(self.success)
[docs]
@dataclass
class PlanState:
r"""Data class representing the state for a motion plan (env-batched).
Tensor fields carry a leading batch dim ``B``: ``qpos:(B, DOF)``,
``xpos:(B, 4, 4)``. Enum/scalar fields are shared across ``B`` (vectorized
envs share the same task skeleton).
"""
move_type: MoveType = MoveType.JOINT_MOVE
"""Type of movement used by the plan."""
move_part: MovePart = MovePart.LEFT
"""Robot part that should move."""
xpos: torch.Tensor | None = None
"""Target TCP pose (Bx4x4) for ``MoveType.EEF_MOVE``."""
qpos: torch.Tensor | None = None
"""Target joint angles for ``MoveType.JOINT_MOVE`` with shape ``(B, DOF)``."""
qvel: torch.Tensor | None = None
"""Target joint velocities for ``MoveType.JOINT_MOVE`` with shape ``(B, DOF)``."""
qacc: torch.Tensor | None = None
"""Target joint accelerations for ``MoveType.JOINT_MOVE`` with shape ``(B, DOF)``."""
is_open: bool = True
"""For ``MoveType.TOOL``, indicates whether to open (``True``) or close (``False``) the tool."""
is_world_coordinate: bool = True
"""``True`` if the target pose is in world coordinates, ``False`` if relative to the current pose."""
pause_seconds: float = 0.0
"""Duration of a pause when ``move_type`` is ``MoveType.PAUSE``."""
[docs]
@classmethod
def from_qpos(
cls,
qpos: torch.Tensor,
*,
move_type: MoveType = MoveType.JOINT_MOVE,
move_part: MovePart = MovePart.LEFT,
**kwargs,
) -> "PlanState":
"""Create a PlanState from batched joint positions ``(B, DOF)``."""
return cls(move_type=move_type, move_part=move_part, qpos=qpos, **kwargs)
[docs]
@classmethod
def from_xpos(
cls,
xpos: torch.Tensor,
*,
move_type: MoveType = MoveType.EEF_MOVE,
move_part: MovePart = MovePart.LEFT,
**kwargs,
) -> "PlanState":
"""Create a PlanState from batched end-effector poses ``(B, 4, 4)``."""
return cls(move_type=move_type, move_part=move_part, xpos=xpos, **kwargs)
[docs]
@classmethod
def single(
cls,
*,
qpos: torch.Tensor | None = None,
xpos: torch.Tensor | None = None,
move_type: MoveType = MoveType.JOINT_MOVE,
move_part: MovePart = MovePart.LEFT,
**kwargs,
) -> "PlanState":
"""B=1 convenience constructor: unsqueezes a single-env qpos/xpos.
Already-batched tensors (2D qpos / 3D xpos) pass through unchanged
(idempotent).
"""
if qpos is not None and qpos.dim() == 1:
qpos = qpos.unsqueeze(0)
if xpos is not None and xpos.dim() == 2:
xpos = xpos.unsqueeze(0)
return cls(
move_type=move_type, move_part=move_part, qpos=qpos, xpos=xpos, **kwargs
)
def interpolate_xpos(
current_xpos: np.ndarray, target_xpos: np.ndarray, num_samples: int
) -> np.ndarray:
"""Interpolate between two poses using vectorized Slerp + linear translation."""
num_samples = max(2, int(num_samples))
interp_ratios = np.linspace(0.0, 1.0, num_samples)
slerp = Slerp(
[0.0, 1.0],
Rotation.from_matrix([current_xpos[:3, :3], target_xpos[:3, :3]]),
)
interp_rots = slerp(interp_ratios).as_matrix()
interp_trans = (1.0 - interp_ratios[:, None]) * current_xpos[:3, 3] + interp_ratios[
:, None
] * target_xpos[:3, 3]
interp_poses = np.repeat(np.eye(4)[None, :, :], num_samples, axis=0)
interp_poses[:, :3, :3] = interp_rots
interp_poses[:, :3, 3] = interp_trans
return interp_poses
def interpolate_xpos_batched(
start_xpos: torch.Tensor, end_xpos: torch.Tensor, num_samples: int
) -> torch.Tensor:
"""Batched pose interpolation.
Args:
start_xpos: Start poses, shape ``(B, 4, 4)``.
end_xpos: End poses, shape ``(B, 4, 4)``.
num_samples: Number of samples to generate (clamped to at least 2).
Returns:
Interpolated poses, shape ``(B, num_samples, 4, 4)``.
"""
num_samples = max(2, int(num_samples))
B = start_xpos.shape[0]
out = torch.eye(4, dtype=start_xpos.dtype, device=start_xpos.device).repeat(
B, num_samples, 1, 1
)
# ``scipy.spatial.transform.Slerp`` interpolates a single path, so we loop
# over envs. B is typically small (number of parallel envs), so this is fine.
for b in range(B):
poses = interpolate_xpos(
start_xpos[b].detach().cpu().numpy(),
end_xpos[b].detach().cpu().numpy(),
num_samples,
)
out[b] = torch.as_tensor(
poses, dtype=start_xpos.dtype, device=start_xpos.device
)
return out
def calculate_point_allocations(
xpos_list: torch.Tensor | np.ndarray,
step_size: float = 0.002,
angle_step: float = np.pi / 90,
device: torch.device = torch.device("cpu"),
) -> List[int]:
"""Calculate interpolation points for each segment with vectorized tensor ops."""
if not isinstance(xpos_list, torch.Tensor):
xpos_tensor = torch.as_tensor(
np.asarray(xpos_list), dtype=torch.float32, device=device
)
else:
xpos_tensor = xpos_list.to(dtype=torch.float32, device=device)
if xpos_tensor.dim() != 3 or xpos_tensor.shape[0] < 2:
return []
start_poses = xpos_tensor[:-1] # [N-1, 4, 4]
end_poses = xpos_tensor[1:] # [N-1, 4, 4]
pos_dists = torch.norm(end_poses[:, :3, 3] - start_poses[:, :3, 3], dim=-1)
pos_points = torch.clamp((pos_dists / step_size).int(), min=1)
rel_rot = torch.matmul(
start_poses[:, :3, :3].transpose(-1, -2), end_poses[:, :3, :3]
)
trace = rel_rot[:, 0, 0] + rel_rot[:, 1, 1] + rel_rot[:, 2, 2]
cos_angle = torch.clamp((trace - 1.0) / 2.0, -1.0 + 1e-6, 1.0 - 1e-6)
angles = torch.acos(cos_angle)
rot_points = torch.clamp((angles / angle_step).int(), min=1)
return torch.maximum(pos_points, rot_points).tolist()