# ----------------------------------------------------------------------------
# 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 typing import Dict, List, Literal, Sequence, Tuple
from dataclasses import dataclass, field
from tensordict import TensorDict
from dexsim.engine import Articulation as _Articulation
from embodichain.lab.sim.cfg import RobotCfg, RobotWorkspaceCfg
from embodichain.lab.sim.motion.solvers import SolverCfg, BaseSolver
from embodichain.lab.sim.objects import Articulation
from embodichain.lab.sim.utility.tensor import to_tensor
from embodichain.lab.sim.motion.workspace.runtime import (
RobotWorkspace,
WorkspaceSample,
)
from embodichain.utils.math import quat_from_matrix, matrix_from_quat
from embodichain.utils.string import (
is_regular_expression,
resolve_matching_names_values,
)
from embodichain.utils import logger
@dataclass
class ControlGroup:
r"""Represents a group of controllable joints in a robot.
Attributes:
joint_names (List[str]): Names of the joints in this control group.
joint_ids (List[int]): IDs corresponding to the joints in this control group.
link_names (List[str]): Names of child links associated with the joints.
"""
joint_names: List[str] = field(default_factory=list)
joint_ids: List[int] = field(default_factory=list)
link_names: List[str] = field(default_factory=list)
def __post_init__(self):
pass
[docs]
class Robot(Articulation):
"""A class representing a batch of robots in the simulation environment.
Robot is a specific type of articulation that can have additional properties or methods.
- `control_parts`: Specify the parts that can be controlled in a different manner. Different part may have
different joint ids, drive properties, pyhsical attributes, kinematic solvers or motion planners.
- `solvers`: Specify the kinematic solvers for the robot.
- `planners`: Specify the motion planner for the robot.
"""
[docs]
def __init__(
self,
cfg: RobotCfg,
entities: List[_Articulation],
device: torch.device = torch.device("cpu"),
) -> None:
self._entities = entities
self.cfg = cfg
# Initialize joint ids for control parts.
self._joint_ids: Dict[str, List[int]] = {}
self._control_groups: Dict[str, ControlGroup] = {}
# Articulation initialization may apply configured qpos limits through
# Robot.set_qpos_limits(), which synchronizes this registry.
self._solvers: Dict[str, BaseSolver] = {}
# Workspaces are loaded lazily so constructing a robot does not perform
# cache I/O unless a task actually requests workspace sampling.
self._workspaces: Dict[str, RobotWorkspace] = {}
if self.cfg.control_parts:
self._init_control_parts(self.cfg.control_parts)
super().__init__(cfg, entities, device)
if self.cfg.solver_cfg:
self.init_solver(self.cfg.solver_cfg)
def __str__(self) -> str:
parent_str = super().__str__()
return (
parent_str
+ f" | control_parts: {self.control_parts}, solvers: {self._solvers}"
)
@property
def control_parts(self) -> Dict[str, List[str]] | None:
"""Get the control parts of the robot."""
return self.cfg.control_parts
[docs]
def attach_workspace(
self, workspace: RobotWorkspace, name: str | None = None
) -> None:
"""Attach a runtime workspace to a control part.
Args:
workspace: Runtime workspace to attach.
name: Control-part name. Use ``None`` for the default solver.
"""
key = name or "default"
if self.control_parts and name is not None and name not in self.control_parts:
raise ValueError(
f"The control part {name!r} does not exist in {self.control_parts}."
)
self._workspaces[key] = workspace.to(self.device)
[docs]
def get_workspace(self, name: str | None = None) -> RobotWorkspace:
"""Get or lazily load the workspace for a control part.
Args:
name: Control-part name. If omitted and exactly one workspace is
configured, that workspace is selected automatically.
Returns:
Runtime robot workspace.
Raises:
ValueError: If no unambiguous workspace configuration exists.
"""
workspace_cfgs = self.cfg.workspace_cfg or {}
key = name or "default"
if name is None and key not in workspace_cfgs and key not in self._workspaces:
available = set(workspace_cfgs) | set(self._workspaces)
if len(available) == 1:
key = next(iter(available))
else:
raise ValueError(
"A control-part name is required when the robot has zero or "
f"multiple workspaces; available workspaces: {sorted(available)}."
)
if key not in self._workspaces:
workspace_cfg = workspace_cfgs.get(key)
if workspace_cfg is None:
raise ValueError(
f"No workspace is configured for control part {key!r}."
)
self._workspaces[key] = RobotWorkspace.from_cache(
workspace_cfg.cache_path,
device=self.device,
voxel_size=workspace_cfg.voxel_size,
)
return self._workspaces[key]
[docs]
def sample_reachable_pose(
self,
name: str | None = None,
env_ids: Sequence[int] | torch.Tensor | None = None,
num_samples: int = 1,
strategy: Literal["point_uniform", "voxel_uniform"] | None = None,
position_bounds: (
torch.Tensor | tuple[Sequence[float], Sequence[float]] | None
) = None,
min_score: float | None = None,
max_attempts: int = 64,
generator: torch.Generator | None = None,
) -> WorkspaceSample:
"""Sample kinematically reachable end-effector poses.
Cached joint configurations are sampled first and then evaluated with
runtime FK. This accounts for the current control-part base pose in
every target environment. A valid sample is kinematically reachable,
but is not guaranteed to have a collision-free trajectory from the
robot's current state.
Args:
name: Control-part name.
env_ids: Target environment IDs. Defaults to all environments.
num_samples: Number of samples per environment.
strategy: Sampling strategy. Defaults to the workspace config.
position_bounds: Optional arena-frame bounds ``((xmin, ymin, zmin),
(xmax, ymax, zmax))`` applied after FK.
min_score: Optional minimum cached reachability score.
max_attempts: Candidate count per environment when applying bounds.
generator: Optional torch random number generator.
Returns:
Batched workspace samples. Entries that could not satisfy bounds
have ``valid=False`` and ``indices=-1``.
Raises:
ValueError: If sampling arguments are invalid.
"""
if num_samples <= 0:
raise ValueError(f"num_samples must be positive; got {num_samples}.")
if max_attempts < num_samples:
raise ValueError(
"max_attempts must be at least num_samples; "
f"got {max_attempts} and {num_samples}."
)
if env_ids is None:
local_env_ids = self._all_indices
elif isinstance(env_ids, slice):
local_env_ids = self._all_indices[env_ids]
else:
local_env_ids = env_ids
if not isinstance(local_env_ids, torch.Tensor):
local_env_ids = torch.as_tensor(
local_env_ids, dtype=torch.long, device=self.device
)
else:
local_env_ids = local_env_ids.to(device=self.device, dtype=torch.long)
batch_size = len(local_env_ids)
workspace = self.get_workspace(name)
workspace_cfgs = self.cfg.workspace_cfg or {}
workspace_key = name or "default"
if name is None and workspace_key not in workspace_cfgs:
configured_keys = set(workspace_cfgs) | set(self._workspaces)
if len(configured_keys) == 1:
workspace_key = next(iter(configured_keys))
fk_name = name
if fk_name is None and workspace_key != "default":
fk_name = workspace_key
workspace_cfg = workspace_cfgs.get(workspace_key)
selected_strategy = strategy or (
workspace_cfg.strategy if workspace_cfg is not None else "voxel_uniform"
)
selected_min_score = (
min_score
if min_score is not None
else (workspace_cfg.min_score if workspace_cfg is not None else None)
)
candidate_count = max_attempts if position_bounds is not None else num_samples
candidate_indices = workspace.sample_indices(
batch_size * candidate_count,
strategy=selected_strategy,
min_score=selected_min_score,
generator=generator,
).reshape(batch_size, candidate_count)
candidate_qpos = workspace.qpos[candidate_indices]
candidate_poses = self.compute_batch_fk(
qpos=candidate_qpos,
name=fk_name,
env_ids=local_env_ids,
to_matrix=True,
)
candidate_valid = torch.ones(
(batch_size, candidate_count), dtype=torch.bool, device=self.device
)
if position_bounds is not None:
bounds = torch.as_tensor(
position_bounds, dtype=torch.float32, device=self.device
)
if bounds.shape != (2, 3):
raise ValueError(
"position_bounds must have shape (2, 3); "
f"got {tuple(bounds.shape)}."
)
positions = candidate_poses[:, :, :3, 3]
candidate_valid = torch.logical_and(
positions >= bounds[0], positions <= bounds[1]
).all(dim=-1)
pose_result = (
torch.eye(4, dtype=candidate_poses.dtype, device=self.device)
.reshape(1, 1, 4, 4)
.repeat(batch_size, num_samples, 1, 1)
)
qpos_result = torch.zeros(
(batch_size, num_samples, workspace.qpos.shape[1]),
dtype=workspace.qpos.dtype,
device=self.device,
)
index_result = torch.full(
(batch_size, num_samples), -1, dtype=torch.long, device=self.device
)
valid_result = torch.zeros(
(batch_size, num_samples), dtype=torch.bool, device=self.device
)
score_result = (
torch.zeros(
(batch_size, num_samples),
dtype=workspace.scores.dtype,
device=self.device,
)
if workspace.scores is not None
else None
)
for batch_index in range(batch_size):
valid_candidates = torch.where(candidate_valid[batch_index])[0][
:num_samples
]
count = len(valid_candidates)
if count == 0:
continue
pose_result[batch_index, :count] = candidate_poses[
batch_index, valid_candidates
]
qpos_result[batch_index, :count] = candidate_qpos[
batch_index, valid_candidates
]
selected_indices = candidate_indices[batch_index, valid_candidates]
index_result[batch_index, :count] = selected_indices
valid_result[batch_index, :count] = True
if score_result is not None:
score_result[batch_index, :count] = workspace.scores[selected_indices]
return WorkspaceSample(
eef_pose=pose_result,
qpos=qpos_result,
indices=index_result,
valid=valid_result,
score=score_result,
)
[docs]
def get_joint_ids(
self, name: str | None = None, remove_mimic: bool = False
) -> List[int]:
"""Get the joint ids of the robot for a specific control part.
Args:
name (str | None): The name of the control part to get the joint ids for. If None, the default part is used.
remove_mimic (bool, optional): If True, mimic joints will be excluded from the returned joint ids. Defaults to False.
Returns:
List[int]: The joint ids of the robot for the specified control part.
"""
if not self.control_parts or name is None:
return (
torch.arange(self.dof, dtype=torch.int32).tolist()
if not remove_mimic
else self.active_joint_ids
)
if name not in self.control_parts:
logger.log_error(
f"The control part '{name}' does not exist in the robot's control parts."
)
return (
self._joint_ids[name]
if not remove_mimic
else [i for i in self._joint_ids[name] if i not in self.mimic_ids]
)
[docs]
def get_link_names(self, name: str | None = None) -> List[str] | None:
"""Get the link names of the robot for a specific control part.
If no control part is specified, return all link names.
Args:
name (str, optional): The name of the control part to get the link names for. If None, the default part is used.
Returns:
List[str]: The link names of the robot for the specified control part.
"""
if not self.control_parts or name is None:
return self.link_names
if name not in self.control_parts:
logger.log_error(
f"The control part '{name}' does not exist in the robot's control parts {self.control_parts}."
)
return self._control_groups[name].link_names
def _resolve_limit_joint_ids(
self,
name: str | None,
joint_ids: Sequence[int] | None,
) -> Sequence[int] | None:
"""Resolve joint selection for robot limit APIs.
When `name` is specified, the control-part joint ids take precedence over
explicit `joint_ids` to preserve the existing Robot setter/getter behavior.
"""
if name is None:
return joint_ids
if not self.control_parts or name not in self.control_parts:
logger.log_error(
f"The control part '{name}' does not exist in the robot's control parts."
)
if joint_ids is not None:
logger.log_warning("`joint_ids` is ignored when `name` is specified.")
return self.get_joint_ids(name=name)
[docs]
def get_qpos_limits(
self,
name: str | None = None,
env_ids: Sequence[int] | torch.Tensor | None = None,
*,
joint_ids: Sequence[int] | torch.Tensor | None = None,
) -> torch.Tensor:
"""Get the joint position limits (qpos) of the robot for a specific control part.
It returns all joint position limits if no control part is specified.
Args:
name (str | None): The name of the control part to get the qpos limits for.
env_ids (Sequence[int] | torch.Tensor | None): The environment ids to get the qpos limits for. If None, all environments are used.
joint_ids (Sequence[int] | torch.Tensor | None): Joint indices to get the qpos limits for.
Must be passed as a keyword argument.
Returns:
torch.Tensor: Joint position limits with shape (N, dof, 2), where N is the number of environments.
"""
resolved_joint_ids = self._resolve_limit_joint_ids(name, joint_ids)
return super().get_qpos_limits(
joint_ids=resolved_joint_ids,
env_ids=env_ids,
)
[docs]
def get_qvel_limits(
self,
name: str | None = None,
env_ids: Sequence[int] | torch.Tensor | None = None,
*,
joint_ids: Sequence[int] | torch.Tensor | None = None,
) -> torch.Tensor:
"""Get the joint velocity limits (qvel) of the robot for a specific control part.
It returns all joint velocity limits if no control part is specified.
Args:
name (str | None): The name of the control part to get the qvel limits for.
env_ids (Sequence[int] | torch.Tensor | None): The environment ids to get the qvel limits for. If None, all environments are used.
joint_ids (Sequence[int] | torch.Tensor | None): Joint indices to get the qvel limits for.
Must be passed as a keyword argument.
Returns:
torch.Tensor: Joint velocity limits with shape (N, dof), where N is the number of environments.
"""
resolved_joint_ids = self._resolve_limit_joint_ids(name, joint_ids)
return super().get_qvel_limits(
joint_ids=resolved_joint_ids,
env_ids=env_ids,
)
[docs]
def get_qf_limits(
self,
name: str | None = None,
env_ids: Sequence[int] | torch.Tensor | None = None,
*,
joint_ids: Sequence[int] | torch.Tensor | None = None,
) -> torch.Tensor:
"""Get the joint effort limits (qf) of the robot for a specific control part.
It returns all joint effort limits if no control part is specified.
Args:
name (str | None): The name of the control part to get the qf limits for.
env_ids (Sequence[int] | torch.Tensor | None): The environment ids to get the qf limits for. If None, all environments are used.
joint_ids (Sequence[int] | torch.Tensor | None): Joint indices to get the qf limits for.
Must be passed as a keyword argument.
Returns:
torch.Tensor: Joint effort limits with shape (N, dof), where N is the number of environments.
"""
resolved_joint_ids = self._resolve_limit_joint_ids(name, joint_ids)
return super().get_qf_limits(
joint_ids=resolved_joint_ids,
env_ids=env_ids,
)
[docs]
def set_qpos_limits(
self,
qpos_limits: torch.Tensor,
name: str | None = None,
env_ids: Sequence[int] | torch.Tensor | None = None,
*,
joint_ids: Sequence[int] | torch.Tensor | None = None,
) -> None:
"""Set the joint position limits (qpos) of the robot.
Args:
qpos_limits: Joint position limits with shape (N, num_joints, 2).
name (str | None): The name of the control part to set the qpos limits for.
env_ids (Sequence[int] | torch.Tensor | None): The environment ids to set the qpos limits for.
joint_ids (Sequence[int] | torch.Tensor | None): Joint indices to set the qpos limits for.
Must be passed as a keyword argument; ignored when ``name`` is provided.
"""
resolved_joint_ids = self._resolve_limit_joint_ids(name, joint_ids)
super().set_qpos_limits(
qpos_limits=qpos_limits,
joint_ids=resolved_joint_ids,
env_ids=env_ids,
)
self._sync_solver_limits(name=name)
[docs]
def set_qvel_limits(
self,
qvel_limits: torch.Tensor,
name: str | None = None,
env_ids: Sequence[int] | torch.Tensor | None = None,
*,
joint_ids: Sequence[int] | torch.Tensor | None = None,
) -> None:
"""Set the joint velocity limits (qvel) of the robot.
Args:
qvel_limits: Joint velocity limits with shape (N, num_joints).
name (str | None): The name of the control part to set the qvel limits for.
env_ids (Sequence[int] | torch.Tensor | None): The environment ids to set the qvel limits for.
joint_ids (Sequence[int] | torch.Tensor | None): Joint indices to set the qvel limits for.
Must be passed as a keyword argument; ignored when ``name`` is provided.
"""
resolved_joint_ids = self._resolve_limit_joint_ids(name, joint_ids)
super().set_qvel_limits(
qvel_limits=qvel_limits,
joint_ids=resolved_joint_ids,
env_ids=env_ids,
)
[docs]
def set_qf_limits(
self,
qf_limits: torch.Tensor,
name: str | None = None,
env_ids: Sequence[int] | torch.Tensor | None = None,
*,
joint_ids: Sequence[int] | torch.Tensor | None = None,
) -> None:
"""Set the joint effort limits (qf) of the robot.
Args:
qf_limits: Joint effort limits with shape (N, num_joints).
name (str | None): The name of the control part to set the qf limits for.
env_ids (Sequence[int] | torch.Tensor | None): The environment ids to set the qf limits for.
joint_ids (Sequence[int] | torch.Tensor | None): Joint indices to set the qf limits for.
Must be passed as a keyword argument; ignored when ``name`` is provided.
"""
resolved_joint_ids = self._resolve_limit_joint_ids(name, joint_ids)
super().set_qf_limits(
qf_limits=qf_limits,
joint_ids=resolved_joint_ids,
env_ids=env_ids,
)
[docs]
def get_proprioception(self) -> TensorDict[str, torch.Tensor]:
"""Gets robot proprioception information, primarily for agent state representation in robot learning scenarios.
The default proprioception information includes:
- qpos: Joint positions.
- qvel: Joint velocities.
- qf: Joint efforts.
Returns:
Dict[str, torch.Tensor]: A dictionary containing the robot's proprioception information
"""
return TensorDict(
qpos=self.body_data.qpos,
qvel=self.body_data.qvel,
qf=self.body_data.qf,
batch_size=[self.num_instances],
device=self.device,
)
[docs]
def set_qpos(
self,
qpos: torch.Tensor,
joint_ids: Sequence[int] | None = None,
env_ids: Sequence[int] | None = None,
target: bool = True,
name: str | None = None,
) -> None:
"""Set the joint positions (qpos) or target positions for the articulation.
Args:
qpos (torch.Tensor): Joint positions with shape (N, dof), where N is the number of environments.
joint_ids (Sequence[int] | None, optional): Joint indices to apply the positions. If None, applies to all joints.
env_ids (Sequence[int] | None): Environment indices to apply the positions. Defaults to all environments.
target (bool): If True, sets target positions for simulation. If False, updates current positions directly.
name (str | None): The name of the control part to set the qpos for. If None, the default part is used.
Raises:
ValueError: If the length of `env_ids` does not match the length of `qpos`.
"""
if name is None:
super().set_qpos(
qpos=qpos,
joint_ids=joint_ids,
env_ids=env_ids,
target=target,
)
else:
if not self.control_parts or name not in self.control_parts:
logger.log_error(
f"The control part '{name}' does not exist in the robot's control parts."
)
part_joint_ids = self.get_joint_ids(name=name)
if joint_ids is not None:
logger.log_warning(f"`joint_ids` is ignored when `name` is specified.")
super().set_qpos(
qpos=qpos,
joint_ids=part_joint_ids,
env_ids=env_ids,
target=target,
)
[docs]
def get_qpos(self, name: str | None = None, target: bool = False) -> torch.Tensor:
"""Get the joint positions (qpos) of the robot.
Args:
name (str | None): The name of the control part to get the qpos for. If None, the default part is used.
target (bool): If True, gets target positions for simulation. If False, gets current positions. The default is False.
Returns:
torch.Tensor: Joint positions with shape (N, dof), where N is the number of environments.
"""
qpos = super().get_qpos(target=target)
if name is None:
return qpos
else:
if not self.control_parts or name not in self.control_parts:
logger.log_error(
f"The control part '{name}' does not exist in the robot's control parts."
)
part_joint_ids = self.get_joint_ids(name=name)
return qpos[:, part_joint_ids]
[docs]
def set_qvel(
self,
qvel: torch.Tensor,
joint_ids: Sequence[int] | None = None,
env_ids: Sequence[int] | None = None,
target: bool = True,
name: str | None = None,
) -> None:
"""Set the joint velocities (qvel) or target velocities for the articulation.
Args:
qvel (torch.Tensor): Joint velocities with shape (N, dof), where N is the number of environments.
joint_ids (Sequence[int] | None, optional): Joint indices to apply the velocities. If None, applies to all joints.
env_ids (Sequence[int] | None): Environment indices to apply the velocities. Defaults to all environments.
target (bool): If True, sets target velocities for simulation. If False, updates current velocities directly.
name (str | None): The name of the control part to set the qvel for. If None, the default part is used.
Raises:
ValueError: If the length of `env_ids` does not match the length of `qvel`.
"""
if name is None:
super().set_qvel(
qvel=qvel,
joint_ids=joint_ids,
env_ids=env_ids,
target=target,
)
else:
if not self.control_parts or name not in self.control_parts:
logger.log_error(
f"The control part '{name}' does not exist in the robot's control parts."
)
part_joint_ids = self.get_joint_ids(name=name)
if joint_ids is not None:
logger.log_warning(f"`joint_ids` is ignored when `name` is specified.")
super().set_qvel(
qvel=qvel,
joint_ids=part_joint_ids,
env_ids=env_ids,
target=target,
)
[docs]
def get_qvel(self, name: str | None = None, target: bool = False) -> torch.Tensor:
"""Get the joint velocities (qvel) of the robot.
Args:
name (str | None): The name of the control part to get the qvel for. If None, the default part is used.
target (bool): If True, gets target velocities for simulation. If False, gets current velocities. The default is False.
Returns:
torch.Tensor: Joint velocities with shape (N, dof), where N is the number of environments.
"""
qvel = super().get_qvel(target=target)
if name is None:
return qvel
else:
if not self.control_parts or name not in self.control_parts:
logger.log_error(
f"The control part '{name}' does not exist in the robot's control parts."
)
part_joint_ids = self.get_joint_ids(name=name)
return qvel[:, part_joint_ids]
[docs]
def set_qf(
self,
qf: torch.Tensor,
joint_ids: Sequence[int] | None = None,
env_ids: Sequence[int] | None = None,
name: str | None = None,
) -> None:
"""Set the joint efforts (qf) for the articulation.
Args:
qf (torch.Tensor): Joint efforts with shape (N, dof), where N is the number of environments.
joint_ids (Sequence[int] | None, optional): Joint indices to apply the efforts. If None, applies to all joints.
env_ids (Sequence[int] | None): Environment indices to apply the efforts. Defaults to all environments.
name (str | None): The name of the control part to set the qf for. If None, the default part is used.
Raises:
ValueError: If the length of `env_ids` does not match the length of `qf`.
"""
if name is None:
super().set_qf(
qf=qf,
joint_ids=joint_ids,
env_ids=env_ids,
)
else:
if not self.control_parts or name not in self.control_parts:
logger.log_error(
f"The control part '{name}' does not exist in the robot's control parts."
)
part_joint_ids = self.get_joint_ids(name=name)
if joint_ids is not None:
logger.log_warning(f"`joint_ids` is ignored when `name` is specified.")
super().set_qf(
qf=qf,
joint_ids=part_joint_ids,
env_ids=env_ids,
)
[docs]
def get_qf(self, name: str | None = None) -> torch.Tensor:
"""Get the joint efforts (qf) of the robot.
Args:
name (str | None): The name of the control part to get the qf for. If None, the default part is used.
Returns:
torch.Tensor: Joint efforts with shape (N, dof), where N is the number of environments.
"""
qf = super().get_qf()
if name is None:
return qf
else:
if not self.control_parts or name not in self.control_parts:
logger.log_error(
f"The control part '{name}' does not exist in the robot's control parts."
)
part_joint_ids = self.get_joint_ids(name=name)
return qf[:, part_joint_ids]
[docs]
def compute_fk(
self,
qpos: torch.Tensor | np.ndarray | None,
name: str | None = None,
link_names: List[str] | None = None,
end_link_name: str | None = None,
root_link_name: str | None = None,
env_ids: Sequence[int] | None = None,
to_matrix: bool = False,
*,
qpos_joint_names: Sequence[str] | None = None,
) -> torch.Tensor:
"""Compute the forward kinematics of the robot given joint positions and optionally a specific part name.
The output pose will be in the local arena frame.
Args:
qpos (torch.Tensor | np.ndarray | None): Joint positions of the robot, (num_envs, num_joints).
name (str | None): The name of the control part to compute the FK for. If None, the default part is used.
link_names (List[str] | None): The names of the links to compute the FK for. If None, all links are used.
end_link_name (str | None): The name of the end link to compute the FK for. If None, the default end link is used.
root_link_name (str | None): The name of the root link to compute the FK for. If None, the default root link is used.
env_ids (Sequence[int] | None): The environment ids to compute the FK for. If None, all environments are used.
to_matrix (bool): If True, returns the transformation in the form of a 4x4 matrix.
qpos_joint_names: Optional names corresponding to the last dimension
of ``qpos`` for full-articulation FK.
Returns:
torch.Tensor: The forward kinematics result with shape (num_envs, 7) or (num_envs, 4, 4) if `to_matrix` is True.
"""
local_env_ids = self._all_indices if env_ids is None else env_ids
if name is None and hasattr(super(), "compute_fk"):
return super().compute_fk(
qpos=qpos,
link_names=link_names,
end_link_name=end_link_name,
root_link_name=root_link_name,
qpos_joint_names=qpos_joint_names,
)
if qpos_joint_names is not None:
raise ValueError(
"qpos_joint_names cannot be combined with named robot control-part FK."
)
if not self._solvers:
logger.log_error(
"No solvers are defined for the robot. Please ensure that the robot has solvers configured."
)
solver = self._solvers.get(name if name is not None else "default", None)
if solver is None:
logger.log_error(
f"The control part '{name}' does not have an associated solver. Please ensure that a valid control part with an available solver is provided."
)
return None
if qpos.dim() == 1:
qpos = qpos.unsqueeze(0)
if qpos.shape[0] != len(local_env_ids):
logger.log_error(
f"Joint positions batch size mismatch. Expected {len(local_env_ids)} but got {qpos.shape[0]}."
)
if qpos.shape[1] != solver.dof:
logger.log_error(
f"Joint positions shape mismatch. Expected {solver.dof} joints, got {qpos.shape[1]}."
)
qpos_ = qpos.to(self.device)
result_matrix = solver.get_fk(qpos=qpos_)
base_pose = self.get_link_pose(
link_name=solver.root_link_name, env_ids=local_env_ids, to_matrix=True
)
result_matrix = torch.bmm(base_pose, result_matrix)
if to_matrix:
return result_matrix
else:
pos = result_matrix[:, :3, 3]
quat = quat_from_matrix(result_matrix[:, :3, :3])
return torch.cat((pos, quat), dim=-1)
[docs]
def compute_ik(
self,
pose: torch.Tensor | np.ndarray,
joint_seed: torch.Tensor | np.ndarray | None = None,
name: str | None = None,
env_ids: Sequence[int] | None = None,
return_all_solutions: bool = False,
) -> Tuple[torch.Tensor, torch.Tensor] | None:
"""Compute the inverse kinematics of the robot given joint positions and optionally a specific part name.
The input pose should be in the local arena frame.
Args:
pose (torch.Tensor): The end effector pose of the robot, (num_envs, 7) or (num_envs, 4, 4).
joint_seed (torch.Tensor | None): The joint positions to use as a seed for the IK computation, (num_envs, dof).
If None, the zero joint positions will be used as the seed.
name (str | None): The name of the control part to compute the IK for. If None, the default part is used.
env_ids (Sequence[int] | None): Environment indices to apply the positions. Defaults to all environments.
return_all_solutions (bool): Whether to return all IK solutions or just the best one. Defaults to False.
Returns:
Tuple[torch.Tensor, torch.Tensor] | None: The success Tensor with shape (num_envs, ) and qpos Tensor with shape (num_envs, max_results, dof), or None if solver not found.
"""
local_env_ids = self._all_indices if env_ids is None else env_ids
solver = self._solvers.get(name if name is not None else "default", None)
if solver is None:
logger.log_error(
f"The control part '{name}' does not have an associated solver. Please ensure that a valid control part with an available solver is provided."
)
return None
pose = to_tensor(pose, device=self.device)
if (pose.dim() == 1 and pose.shape[1] == 7) or (
pose.dim() == 2 and pose.shape[1] == 4
):
pose = pose.unsqueeze(0)
if pose.shape[0] != len(local_env_ids):
logger.log_error(
f"Pose batch size mismatch. Expected {len(local_env_ids)} but got {pose.shape[0]}."
)
if joint_seed is not None:
joint_seed = to_tensor(joint_seed, device=self.device)
if joint_seed.dim() == 1:
joint_seed = joint_seed.unsqueeze(0)
if joint_seed.shape[0] != len(local_env_ids):
logger.log_error(
f"Joint seed batch size mismatch. Expected {len(local_env_ids)} but got {joint_seed.shape[0]}."
)
if pose.shape[-1] == 7 and pose.dim() == 2:
# Convert pose from (batch, 7) to (batch, 4, 4)
pos = pose[:, :3]
quat = pose[:, 3:]
# Convert quaternion to rotation matrix
rot = matrix_from_quat(quat)
# Build homogeneous transformation matrix efficiently
pose = torch.eye(4, device=pose.device).repeat(pose.shape[0], 1, 1)
pose[:, :3, :3] = rot
pose[:, :3, 3] = pos
base_pose = self.get_link_pose(
link_name=solver.root_link_name, env_ids=local_env_ids, to_matrix=True
)
pose = torch.bmm(torch.inverse(base_pose), pose)
ret, qpos = solver.get_ik(
target_xpos=pose,
qpos_seed=joint_seed,
return_all_solutions=return_all_solutions,
)
dof = qpos.shape[-1]
if not return_all_solutions:
qpos = qpos.reshape(-1, dof)
return ret.to(self.device), qpos.to(self.device)
[docs]
def compute_batch_fk(
self,
qpos: torch.Tensor,
name: str,
env_ids: Sequence[int] | None = None,
to_matrix: bool = False,
):
"""Compute the forward kinematics of the robot given joint positions and optionally a specific part name.
The output pose will be in the local arena frame.
Args:
qpos (torch.Tensor | np.ndarray | None): Joint positions of the robot, (num_envs, n_batch, num_joints).
name (str | None): The name of the control part to compute the FK for. If None, the default part is used.
env_ids (Sequence[int] | None): The environment ids to compute the FK for. If None, all environments are used.
to_matrix (bool): If True, returns the transformation in the form of a 4x4 matrix.
Returns:
torch.Tensor: The forward kinematics result with shape (num_envs, batch, 7) or (num_envs, batch, 4, 4) if `to_matrix` is True.
"""
local_env_ids = self._all_indices if env_ids is None else env_ids
if not self._solvers:
logger.log_error(
"No solvers are defined for the robot. Please ensure that the robot has solvers configured."
)
solver = self._solvers.get(name if name is not None else "default", None)
if solver is None:
logger.log_error(
f"The control part '{name}' does not have an associated solver. Please ensure that a valid control part with an available solver is provided."
)
return None
if qpos.shape[0] != len(local_env_ids):
logger.log_error(
f"Joint positions batch size mismatch. Expected {len(local_env_ids)} but got {qpos.shape[0]}."
)
if qpos.shape[2] != solver.dof:
logger.log_error(
f"Joint positions shape mismatch. Expected {solver.dof} joints, got {qpos.shape[1]}."
)
qpos_ = qpos.to(self.device)
n_batch = qpos_.shape[1]
qpos_batch = qpos_.reshape(-1, solver.dof)
xpos_batch = solver.get_fk(qpos=qpos_batch)
# get xpos from link root
base_xpos_n_envs = self.get_link_pose(
link_name=solver.root_link_name, env_ids=local_env_ids, to_matrix=True
)
base_xpos_batch = (
base_xpos_n_envs[:, None, :, :].repeat(1, n_batch, 1, 1).reshape(-1, 4, 4)
)
result_matrix = torch.bmm(base_xpos_batch, xpos_batch)
if to_matrix:
result_matrix = result_matrix.reshape(len(local_env_ids), n_batch, 4, 4)
return result_matrix
else:
pos = result_matrix[:, :3, 3]
quat = quat_from_matrix(result_matrix[:, :3, :3])
result = torch.cat((pos, quat), dim=-1)
result = result.reshape(len(local_env_ids), n_batch, 7)
return result
[docs]
def compute_batch_ik(
self,
pose: torch.Tensor | np.ndarray,
joint_seed: torch.Tensor | np.ndarray | None,
name: str,
env_ids: Sequence[int] | None = None,
):
"""Compute the inverse kinematics of the robot given joint positions and optionally a specific part name.
The input pose should be in the local arena frame.
Args:
pose (torch.Tensor): The end effector pose of the robot, (num_envs, n_batch, 7) or (num_envs, n_batch, 4, 4).
joint_seed (torch.Tensor | None): The joint positions to use as a seed for the IK computation, (num_envs, n_batch, dof). If None, the zero joint positions will be used as the seed.
name (str | None): The name of the control part to compute the IK for. If None, the default part is used.
env_ids (Sequence[int] | None): Environment indices to apply the positions. Defaults to all environments.
Returns:
Tuple[torch.Tensor, torch.Tensor]:
Success Tensor with shape (num_envs, n_batch)
Qpos Tensor with shape (num_envs, n_batch, dof).
"""
local_env_ids = self._all_indices if env_ids is None else env_ids
solver = self._solvers.get(name if name is not None else "default", None)
if solver is None:
logger.log_error(
f"The control part '{name}' does not have an associated solver. Please ensure that a valid control part with an available solver is provided."
)
return None
pose = to_tensor(pose, device=self.device)
if pose.shape[0] != len(local_env_ids):
logger.log_error(
f"Pose batch size mismatch. Expected {len(local_env_ids)} but got {pose.shape[0]}."
)
n_batch = pose.shape[1]
n_dof = solver.dof
if joint_seed is None:
joint_seed = torch.zeros(
(len(local_env_ids), n_batch, n_dof),
dtype=torch.float32,
device=self.device,
)
if joint_seed.shape[0] != len(local_env_ids):
logger.log_error(
f"Joint seed env size mismatch. Expected {len(local_env_ids)} but got {joint_seed.shape[0]}."
)
if joint_seed.shape[1] != n_batch:
logger.log_error(
f"Joint seed batch size mismatch. Expected {n_batch} but got {joint_seed.shape[1]}."
)
if joint_seed.shape[-1] != n_dof:
logger.log_error(
f"Joint seed dof size mismatch. Expected {n_batch} but got {joint_seed.shape[-1]}."
)
if pose.shape[-1] == 7 and pose.dim() == 3:
# Convert pose from (num_envs, n_batch, 7) to (num_envs * n_batch, 4, 4)
pose_batch = pose.reshape(-1, 7)
pos = pose_batch[:, :3]
quat = pose_batch[:, 3:]
# Convert quaternion to rotation matrix
rot = matrix_from_quat(quat)
# Build homogeneous transformation matrix efficiently
pose_batch = torch.eye(4, device=pose.device).repeat(
pose_batch.shape[0], 1, 1
)
pose_batch[:, :3, :3] = rot
pose_batch[:, :3, 3] = pos
else:
# Convert pose from (num_envs, n_batch, 4, 4) to (num_envs * n_batch, 4, 4)
pose_batch = pose.reshape(-1, 4, 4)
# get xpos from link root
base_xpos_n_envs = self.get_link_pose(
link_name=solver.root_link_name, env_ids=local_env_ids, to_matrix=True
)
base_inv_xpos_n_envs = torch.inverse(base_xpos_n_envs)
base_inv_xpos_batch = (
base_inv_xpos_n_envs[:, None, :, :]
.repeat(1, n_batch, 1, 1)
.reshape(-1, 4, 4)
)
pose_batch = torch.bmm(base_inv_xpos_batch, pose_batch)
joint_seed_batch = joint_seed.reshape(-1, n_dof)
ret, qpos_batch = solver.get_ik(
target_xpos=pose_batch,
qpos_seed=joint_seed_batch,
return_all_solutions=False,
)
ret = ret.reshape(len(local_env_ids), n_batch)
qpos = qpos_batch.reshape(len(local_env_ids), n_batch, n_dof)
return ret, qpos
def _init_control_parts(self, control_parts: Dict[str, List[str]]) -> None:
"""Initialize the control parts of the robot.
Args:
control_parts (Dict[str, List[str]]): A dictionary where keys are control part names and values are lists of
joint names or regular expressions that match joint names.
"""
joint_name_to_ids = {
name: i
for i, name in enumerate(self._entities[0].get_actived_joint_names())
}
for name, joint_names in control_parts.items():
# convert joint_names which is a regular expression to a list of joint names
joint_names_expanded = []
for jn in joint_names:
if is_regular_expression(jn):
_, names, _ = resolve_matching_names_values(
{jn: None}, self.joint_names
)
joint_names_expanded.extend(names)
else:
joint_names_expanded.append(jn)
self._joint_ids[name] = [
joint_name_to_ids[joint_name]
for joint_name in joint_names_expanded
if joint_name in joint_name_to_ids
]
if len(self._joint_ids[name]) != len(joint_names_expanded):
logger.log_error(
f"joint names in control part '{name}' do not match the robot's joint names. The full joint names are: {self.joint_names}."
)
self.cfg.control_parts[name] = joint_names_expanded
# Initialize control groups
self._control_groups = self._extract_control_groups()
[docs]
def set_joint_drive(
self,
stiffness: torch.Tensor | None = None,
damping: torch.Tensor | None = None,
max_effort: torch.Tensor | None = None,
max_velocity: torch.Tensor | None = None,
friction: torch.Tensor | None = None,
armature: torch.Tensor | None = None,
drive_type: str = "force",
joint_ids: Sequence[int] | None = None,
env_ids: Sequence[int] | None = None,
) -> None:
"""Set the drive properties for the robot.
Different from Articulation, default drive type is 'force' instead of 'none'
Args:
stiffness (torch.Tensor): The stiffness of the joint drive with shape (len(env_ids), len(joint_ids)).
damping (torch.Tensor): The damping of the joint drive with shape (len(env_ids), len(joint_ids)).
max_effort (torch.Tensor): The maximum effort of the joint drive with shape (len(env_ids), len(joint_ids)).
max_velocity (torch.Tensor): The maximum velocity of the joint drive with shape (len(env_ids), len(joint_ids)).
friction (torch.Tensor): The joint friction coefficient with shape (len(env_ids), len(joint_ids)).
armature (torch.Tensor): The joint armature with shape (len(env_ids), len(joint_ids)).
drive_type (str, optional): The type of drive to apply. Defaults to "force".
joint_ids (Sequence[int] | None, optional): The joint indices to apply the drive to. If None, applies to all joints. Defaults to None.
env_ids (Sequence[int] | None, optional): The environment indices to apply the drive to. If None, applies to all environments. Defaults to None.
"""
super().set_joint_drive(
stiffness=stiffness,
damping=damping,
max_effort=max_effort,
max_velocity=max_velocity,
friction=friction,
armature=armature,
drive_type=drive_type,
joint_ids=joint_ids,
env_ids=env_ids,
)
def _set_default_joint_drive(self) -> None:
"""Set default joint drive parameters based on the configuration."""
import numbers
from embodichain.utils.string import resolve_matching_names_values
drive_props = [
("damping", self.default_joint_damping),
("stiffness", self.default_joint_stiffness),
("max_effort", self.default_joint_max_effort),
("max_velocity", self.default_joint_max_velocity),
("friction", self.default_joint_friction),
("armature", self.default_joint_armature),
]
for prop_name, default_array in drive_props:
value = getattr(self.cfg.drive_pros, prop_name, None)
if value is None:
continue
if isinstance(value, numbers.Number):
default_array[:] = value
else:
try:
control_part_dict = {}
value_copy = value.copy()
if self.control_parts:
# Extract control part and map the corresponding joint names
for key in value.keys():
if key in self.control_parts:
control_part_dict[key] = value_copy.pop(key)
indices, _, values = resolve_matching_names_values(
value_copy, self.joint_names
)
if self.control_parts:
# Add control part joints to indices and values
for part_name, part_value in control_part_dict.items():
part_joint_names = self.control_parts[part_name]
(
part_indices,
_,
part_values,
) = resolve_matching_names_values(
{jn: part_value for jn in part_joint_names},
self.joint_names,
)
indices.extend(part_indices)
values.extend(part_values)
default_array[:, indices] = torch.as_tensor(
values, dtype=torch.float32, device=self.device
)
except Exception as e:
logger.log_error(f"Failed to set {prop_name}: {e}")
drive_pros = self.cfg.drive_pros
if isinstance(drive_pros, dict):
drive_type = drive_pros.get("drive_type", "force")
else:
drive_type = getattr(drive_pros, "drive_type", "force")
# Apply drive parameters to all articulations in the batch
self.set_joint_drive(
stiffness=self.default_joint_stiffness,
damping=self.default_joint_damping,
max_effort=self.default_joint_max_effort,
max_velocity=self.default_joint_max_velocity,
friction=self.default_joint_friction,
armature=self.default_joint_armature,
drive_type=drive_type,
)
def _sync_solver_limits(self, name: str | None = None) -> None:
"""Synchronize solver joint limits with the robot's effective qpos limits."""
if not self._solvers:
return
if not self.control_parts:
solver = self._solvers.get("default")
if solver is not None:
solver.update_with_robot_limit(self._data.qpos_limits[0])
return
part_names = [name] if name is not None else list(self.control_parts.keys())
for part_name in part_names:
solver = self._solvers.get(part_name)
if solver is None:
continue
joint_ids = self.get_joint_ids(name=part_name)
joint_limits = self._data.qpos_limits[0][joint_ids]
solver.update_with_robot_limit(joint_limits)
[docs]
def init_solver(self, cfg: SolverCfg | Dict[str, SolverCfg]) -> None:
"""Initialize the kinematic solver for the robot.
Args:
cfg (SolverCfg | Dict[str, SolverCfg]): The configuration for the kinematic solver.
"""
self.cfg: RobotCfg
if isinstance(cfg, SolverCfg):
if self.control_parts:
logger.log_error(
"Control parts are defined in the robot configuration, solver_cfg must be a dictionary."
)
if cfg.urdf_path is None:
cfg.urdf_path = self.cfg.fpath
self._solvers["default"] = cfg.init_solver(device=self.device)
self._sync_solver_limits()
elif isinstance(cfg, Dict):
if isinstance(self.cfg.control_parts, Dict) is False:
logger.log_error(
"When `solver_cfg` is a dictionary, `control_parts` must also be a dictionary."
)
# If solver_cfg is a dictionary, iterate through it to create solvers
for name, solver_cfg in cfg.items():
if solver_cfg.urdf_path is None:
solver_cfg.urdf_path = self.cfg.fpath
_, part_names, value = resolve_matching_names_values(
{name: solver_cfg}, self.cfg.control_parts.keys()
)
for part_name in part_names:
if (
not hasattr(solver_cfg, "joint_names")
or solver_cfg.joint_names is None
):
solver_cfg.joint_names = self.cfg.control_parts[part_name]
self._solvers[part_name] = solver_cfg.init_solver(
device=self.device
)
self._sync_solver_limits()
[docs]
def get_solver(self, name: str | None = None) -> BaseSolver | None:
"""Get the kinematic solver for a specific control part.
Args:
name (str | None): The name of the control part to get the solver for. If None, the default part is used.
Returns:
BaseSolver | None: The kinematic solver for the specified control part, or None if not found.
"""
if not self._solvers:
logger.log_error(
"No solvers are defined for the robot. Please ensure that the robot has solvers configured."
)
return None
return self._solvers.get(name if name is not None else "default", None)
[docs]
def get_control_part_base_pose(
self,
name: str | None = None,
env_ids: Sequence[int] | None = None,
to_matrix: bool = False,
) -> torch.Tensor:
"""Retrieves the base pose of the control part for a specified robot.
Args:
name (str | None): The name of the control part the solver adhere to. If None, the default solver is used.
env_ids (Sequence[int] | None): A sequence of environment IDs to specify the environments. If None, all indices are used.
to_matrix (bool): If True, returns the pose in the form of a 4x4 matrix.
Returns:
The pose of the specified link in the form of a matrix.
"""
local_env_ids = self._all_indices if env_ids is None else env_ids
root_link_name = None
if name in self._control_groups:
root_link_name = self._control_groups[name].link_names[0]
return self.get_link_pose(
link_name=root_link_name, env_ids=local_env_ids, to_matrix=to_matrix
)
[docs]
def get_control_part_link_names(self, name: str | None = None) -> List[str]:
"""Get the link names of the control part.
Args:
name (str | None): The name of the control part. If None, return all link names.
Returns:
List[str]: link names of the control part.
"""
if name is None:
return self.link_names
if name in self._control_groups:
return self._control_groups[name].link_names
else:
logger.log_warning(
f"The control part '{name}' does not exist in the robot's control parts."
)
return []
def _extract_control_groups(self) -> Dict[str, ControlGroup]:
r"""Extract control groups from the active joint names.
This method creates a dictionary of control groups where each control
group is associated with its corresponding joint names. It utilizes
the `_extract_control_group` method to populate the control groups.
Returns:
Dict[str, ControlGroup]: A dictionary mapping control group names
to their corresponding ControlGroup instances.
"""
if not self.control_parts:
return {}
control_groups = {
control_group_name: self._extract_control_group(joint_names)
for control_group_name, joint_names in self.control_parts.items()
}
return control_groups
def _extract_control_group(self, joint_names: List[str]) -> ControlGroup:
r"""Extract a control group from the given list of joint names.
Args:
joint_names (List[str]): A list of joint names
to be included in the control group.
Returns:
ControlGroup: An instance of ControlGroup containing the specified joints
and their associated links.
"""
control_group = ControlGroup()
joint_id_list = []
for joint_name in joint_names:
if joint_name in self.joint_names:
joint_index = self.joint_names.index(joint_name)
joint_id_list.append(joint_index)
control_group.joint_names.append(joint_name)
# Set root link for first joint
if len(control_group.link_names) == 0:
parent_names = self._entities[0].get_ancestral_link_names(
joint_index
)
control_group.link_names.extend(parent_names)
child_name = self._entities[0].get_child_link_name(joint_index)
control_group.link_names.append(child_name)
control_group.joint_ids = joint_id_list
return control_group
[docs]
def build_pk_serial_chain(self) -> None:
"""Build the kinematic serial chain for the robot.
This method is mainly used for robot learning scenarios, for example:
- Imitation learning dataset generation.
"""
self.pk_serial_chain = self.cfg.build_pk_serial_chain(device=self.device)
[docs]
def set_physical_visible(
self,
visible: bool = True,
control_part: str | None = None,
rgba: Sequence[float] | None = None,
):
"""set collision of the robot or a specific control part.
Args:
visible (bool, optional): is collision body visible. Defaults to True.
control_part (str | None, optional): control part to set visibility. Defaults to None. If None, all links are set.
rgba (Sequence[float] | None, optional): collision body visible rgba. It will be defined at the first time the function is called. Defaults to None.
"""
rgba = rgba if rgba is not None else (0.8, 0.2, 0.2, 0.7)
if len(rgba) != 4:
logger.log_error(f"Invalid rgba {rgba}, should be a sequence of 4 floats.")
rgba = np.array(
[
rgba[0],
rgba[1],
rgba[2],
rgba[3],
]
)
link_names = self.get_control_part_link_names(name=control_part)
# create collision visible node if not exist
if visible:
for i, env_idx in enumerate(self._all_indices):
for link_name in link_names:
if self._has_collision_visible_node_dict[link_name] is False:
self._entities[env_idx].create_physical_visible_node(
rgba, link_name
)
self._has_collision_visible_node_dict[link_name] = True
# set visibility
for i, env_idx in enumerate(self._all_indices):
for link_name in link_names:
self._entities[env_idx].set_physical_visible(visible, link_name)
[docs]
def destroy(self) -> None:
return super().destroy()
__all__ = ["ControlGroup", "Robot"]