# ----------------------------------------------------------------------------
# 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 TYPE_CHECKING, Dict, List, Literal, Sequence, Tuple
from dataclasses import dataclass, field
from tensordict import TensorDict
from embodichain.lab.sim.cfg import RobotCfg, RobotWorkspaceCfg
from embodichain.lab.sim.cfg.source import _asset_file_path
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
if TYPE_CHECKING:
from dexsim.scene import SpawnedArticulation
@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,
device: torch.device = torch.device("cpu"),
) -> None:
"""Create an unregistered robot facade."""
# 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] = {}
super().__init__(cfg, device)
def __str__(self) -> str:
parent_str = super().__str__()
return (
parent_str
+ f" | control_parts: {self.control_parts}, solvers: {self._solvers}"
)
[docs]
def attach_spawn_handles(
self,
entities: Sequence[SpawnedArticulation],
) -> None:
"""Store handles and expose robot metadata without creating Batch data.
Runtime Batch/Data initialization remains the responsibility of
``bind_spawn()`` after Spawn finalization.
"""
super().attach_spawn_handles(entities)
if self.cfg.control_parts:
self._init_control_parts(self.cfg.control_parts)
def _initialize_spawn_bound_extension(self) -> None:
"""Initialize robot-specific runtime state after Scene binding."""
if self.cfg.control_parts:
self._init_control_parts(self.cfg.control_parts)
if self.cfg.solver_cfg:
self.init_solver(self.cfg.solver_cfg)
@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)`` in ``(x, y, z, qx, qy, qz, qw)`` order, 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 as ``(num_envs, 7)`` in
``(x, y, z, qx, qy, qz, qw)`` order 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)`` in ``xyz + xyzw`` order, 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]
xpos_batch = solver.get_fk_batch(qpos_)
# 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
)
result_matrix = torch.matmul(base_xpos_n_envs[:, None], xpos_batch)
result_matrix = result_matrix.reshape(-1, 4, 4)
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,
*,
continuous: bool = False,
) -> tuple[torch.Tensor, torch.Tensor] | None:
"""Compute batched inverse kinematics for arena-frame poses.
When ``continuous`` is enabled, all solver candidates are generated
in one solver call and the branch nearest to the previous path sample
is selected sequentially. The selected solver must support continuous
candidate selection.
Args:
pose: End-effector poses shaped ``(B, N, 7)`` or
``(B, N, 4, 4)`` in the local arena frame.
joint_seed: Independent seeds shaped ``(B, N, DOF)``. In
continuous mode, the initial path seed shaped ``(B, DOF)``.
Defaults to zero independent seeds when omitted.
name: Control part whose solver should be used.
env_ids: Optional environment indices. Defaults to all
environments.
continuous: Preserve one temporally continuous IK branch across
the ``N`` path samples. Defaults to ``False``.
Returns:
Per-sample success shaped ``(B, N)`` and joint positions shaped
``(B, N, DOF)``, or ``None`` when no solver is configured.
Raises:
ValueError: If an input shape is invalid or continuous selection
is unsupported by the configured solver.
"""
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
if continuous and not solver.supports_continuous_batch_ik:
raise ValueError(
f"Solver for {name!r} does not support continuous batch IK."
)
pose = to_tensor(pose, device=self.device)
batch_size = len(local_env_ids)
if pose.ndim not in {3, 4} or pose.shape[0] != batch_size:
raise ValueError(
f"pose must have shape ({batch_size}, N, 7) or "
f"({batch_size}, N, 4, 4)."
)
n_batch = pose.shape[1]
n_dof = solver.dof
if continuous:
if joint_seed is None:
joint_seed_tensor = torch.zeros(
(batch_size, n_dof),
dtype=torch.float32,
device=self.device,
)
else:
joint_seed_tensor = to_tensor(joint_seed, device=self.device)
if joint_seed_tensor.shape != (batch_size, n_dof):
raise ValueError(
f"joint_seed must have shape ({batch_size}, {n_dof}) "
"when continuous=True."
)
elif joint_seed is None:
joint_seed = torch.zeros(
(batch_size, n_batch, n_dof),
dtype=torch.float32,
device=self.device,
)
else:
joint_seed = to_tensor(joint_seed, device=self.device)
if not continuous and joint_seed.shape != (batch_size, n_batch, n_dof):
raise ValueError(
f"joint_seed must have shape ({batch_size}, {n_batch}, {n_dof})."
)
if pose.shape[-1] == 7 and pose.ndim == 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
elif pose.shape[-2:] == (4, 4) and pose.ndim == 4:
# Convert pose from (num_envs, n_batch, 4, 4) to (num_envs * n_batch, 4, 4)
pose_batch = pose.reshape(-1, 4, 4)
else:
raise ValueError(
f"pose must have shape ({batch_size}, N, 7) or "
f"({batch_size}, N, 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)
pose_batch = torch.matmul(
base_inv_xpos_n_envs[:, None],
pose_batch.reshape(len(local_env_ids), n_batch, 4, 4),
)
if continuous:
repeated_seed = (
joint_seed_tensor[:, None]
.expand(batch_size, n_batch, n_dof)
.reshape(-1, n_dof)
)
candidate_valid, candidate_qpos = solver.get_ik(
target_xpos=pose_batch.reshape(-1, 4, 4),
qpos_seed=repeated_seed,
return_all_solutions=True,
)
return solver._select_continuous_ik_path(
candidate_qpos.reshape(batch_size, n_batch, -1, n_dof),
candidate_valid.reshape(batch_size, n_batch, -1),
joint_seed_tensor,
)
joint_seed_batch = joint_seed.reshape(-1, n_dof)
ret, qpos_batch = solver.get_ik(
target_xpos=pose_batch.reshape(-1, 4, 4),
qpos_seed=joint_seed_batch,
return_all_solutions=False,
)
ret = ret.reshape(batch_size, n_batch)
qpos = qpos_batch.reshape(batch_size, 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._state_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 | None = "force",
joint_ids: Sequence[int] | None = None,
env_ids: Sequence[int] | None = None,
*,
target_mode: str | int | None = None,
) -> None:
"""Set the drive properties for the robot.
With no explicit mode, robots retain their position+velocity force
drive default.
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: Drive type 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.
target_mode: Portable target mode name or integer value 0 through 4.
"""
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,
target_mode=target_mode,
)
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
joint_property_targets = [
("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 joint_property_targets:
value = getattr(self.cfg.joint_drive_props, 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}")
joint_drive_props = self.cfg.joint_drive_props
if isinstance(joint_drive_props, dict):
drive_type = joint_drive_props.get("drive_type")
target_mode = joint_drive_props.get("target_mode")
else:
drive_type = getattr(joint_drive_props, "drive_type", None)
target_mode = getattr(joint_drive_props, "target_mode", None)
if isinstance(target_mode, dict):
logger.log_warning(
"Per-joint target_mode mappings require a Spawn-bound robot; "
"the retained raw-robot path preserves its current target modes."
)
target_mode = None
# 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,
target_mode=target_mode,
)
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 = _asset_file_path(self.cfg)
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 = _asset_file_path(self.cfg)
_, 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 = []
state_joint_ids = {
name: index for index, name in enumerate(self._state_joint_names())
}
source_joint_ids = {
name: index
for index, name in enumerate(self._entities[0].get_actived_joint_names())
}
for joint_name in joint_names:
if joint_name in state_joint_ids and joint_name in source_joint_ids:
joint_id_list.append(state_joint_ids[joint_name])
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(
source_joint_ids[joint_name]
)
control_group.link_names.extend(parent_names)
child_name = self._entities[0].get_child_link_name(
source_joint_ids[joint_name]
)
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)
if self.is_spawn_bound:
for env_idx in self._all_indices:
entity = self._entities[env_idx]
for link_name in link_names:
self._spawn_result.set_physical_visible(
(entity, link_name), rgba, visible
)
for link_name in link_names:
self._has_collision_visible_node_dict[link_name] = True
return
# 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"]