URScriptPrimitiveTransform#

class torchrl.envs.transforms.URScriptPrimitiveTransform(*args: Any, execute: bool = False, multi_action_dim: int = 1, stack_rewards: bool = True, stack_observations: bool = False, **kwargs: Any)[source]#

URScript-style preset of MacroPrimitiveTransform.

This specialization is scoped to six-joint UR-style arms with a scalar gripper command. The policy-facing action is a RobotMacroAction placed under action_key. The transform reads the arm and gripper joint observations to build the interpolation start, maps each primitive to a seven-dimensional joint-position + gripper destination (running Cartesian inverse kinematics for reach_pose), and delegates fixed-length interpolation / execution to the generic base.

Parameters:

execute – if True, return Compose(MultiAction(...), transform).
action_key – key carrying the macro action and the expanded low-level sequence.
robot_qpos_key – observation key for the six arm joints.
gripper_qpos_key – observation key for the gripper joints.
macro_steps – interpolated low-level actions per primitive.
settle_steps – repeated final actions appended per primitive.
action_dim – low-level action dimension (six joints + one gripper).
cartesian_solver – optional callable mapping (target_pose, start_action) to a low-level action. When omitted, the transform uses a parent env’s _cartesian_pose_to_joint_target hook.
open_gripper_ctrl – low-level gripper command for an open gripper.
close_gripper_ctrl – low-level gripper command for a closed gripper.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> from torchrl.envs import RobotMacroAction, URScriptPrimitiveTransform
>>> td = TensorDict({
...     "action": RobotMacroAction.reach_joints(joints=torch.ones(1, 6), steps=2),
...     "robot_qpos": torch.zeros(1, 6),
...     "gripper_qpos": torch.zeros(1, 2),
... }, batch_size=[1])
>>> URScriptPrimitiveTransform().inv(td)["action"].shape
torch.Size([1, 2, 7])

action_sequence(tensordict: TensorDictBase, primitive_id: int | IntEnum | None = None, *, target_pose: Tensor | None = None, target_qpos: Tensor | None = None, gripper: float | Tensor | None = None, steps: int | None = None, settle_steps: int | None = None) → Tensor[source]#: Expand a UR primitive into its low-level sequence without executing.

current_action(tensordict: TensorDictBase, batch_shape: Size, device: device, dtype: dtype, action_dim: int) → Tensor[source]#

Return the low-level action used as the interpolation start.

The base implementation starts every macro from the zero action: in the inverse path action_key carries the incoming macro action (the target), so it must not be read back here as the start. Subclasses that can read the controlled state from observations (e.g. joint positions) override this hook.

low_level_action(robot_qpos: Tensor, gripper: float | Tensor | None = None) → Tensor[source]#: Build a low-level joint-position + gripper action.

Return a copy of tensordict carrying one RobotMacroAction.

Maps a URScript primitive id (and its pose / joint / gripper arguments) onto a RobotMacroAction placed under action_key.

primitive_enum#: alias of URScriptPrimitive

transform_input_spec(input_spec: Composite) → Composite[source]#

Transforms the input spec such that the resulting spec matches transform mapping.

Parameters:: input_spec (TensorSpec) – spec before the transform
Returns:: expected spec after the transform

URScriptPrimitiveTransform#

Docs

Tutorials

Resources