diff --git a/agent_context/topics/motion-planning/motion-planning.md b/agent_context/topics/motion-planning/motion-planning.md index e6028a532..608ff68de 100644 --- a/agent_context/topics/motion-planning/motion-planning.md +++ b/agent_context/topics/motion-planning/motion-planning.md @@ -9,7 +9,6 @@ | TOPPRA planner | `embodichain/lab/sim/planners/toppra_planner.py` → `ToppraPlanner`, `ToppraPlannerCfg`, `ToppraPlanOptions` | | Neural planner | `embodichain/lab/sim/planners/neural_planner.py` → `NeuralPlanner`, `NeuralPlannerCfg`, `NeuralPlanOptions` | | cuRobo planner | `embodichain/lab/sim/planners/curobo/curobo_planner.py` → `CuroboPlanner`, `CuroboPlannerCfg`, `CuroboWorldCfg`, `CuroboPlanOptions` | -| Planner assets | `embodichain/data/assets/planner_assets.py` → `download_neural_planner_checkpoint()` | | Motion generator | `embodichain/lab/sim/planners/motion_generator.py` → `MotionGenerator`, `MotionGenCfg`, `MotionGenOptions` | | Planner utilities & data types | `embodichain/lab/sim/planners/utils.py` → `PlanState`, `PlanResult`, `MoveType`, `MovePart`, `TrajectorySampleMethod`, `interpolate_xpos_batched` | @@ -40,7 +39,7 @@ Config hierarchy: ``` BasePlannerCfg robot_uid (MISSING), planner_type ├─ ToppraPlannerCfg planner_type = "toppra", max_workers, mp_context - └─ NeuralPlannerCfg planner_type = "neural", checkpoint_path (MISSING) + └─ NeuralPlannerCfg planner_type = "neural", onnx_model_path (MISSING) MotionGenCfg planner_cfg (MISSING — must be a BasePlannerCfg subclass) @@ -89,11 +88,11 @@ Worker details: Learning-based EEF waypoint planner. Franka Panda only. -- Checkpoint: `download_neural_planner_checkpoint()` from HuggingFace (gated, needs `HF_TOKEN`) +- Runtime: standalone `.onnx` policy with normalization embedded; install the `nmg` extra - Use via `MotionGenerator` with `planner_type="neural"` and `plan_opts=NeuralPlanOptions(...)` - Input: `EEF_MOVE` `PlanState` list with batched `xpos:(B, 4, 4)` -- Key cfg: `checkpoint_path` (from download), `control_part` -- Natively batched: transformer forward, reach checks, and convergence holds all operate on `(B, ...)`. +- Key cfg: `onnx_model_path`, `control_part`, `num_waypoints`, `policy_frame_from_world`, `runtime_tcp_from_policy_tcp` +- The NMG exporter produces a dynamic-batch ONNX policy; NeuralPlanner rolls out all environments together. ### CuroboPlanner collision worlds diff --git a/docs/source/overview/sim/planners/index.rst b/docs/source/overview/sim/planners/index.rst index 3ee955363..9ea901a09 100644 --- a/docs/source/overview/sim/planners/index.rst +++ b/docs/source/overview/sim/planners/index.rst @@ -27,8 +27,8 @@ The `embodichain` project provides a unified interface for robot trajectory plan These tools can be used to generate smooth and dynamically feasible robot trajectories. Install NVIDIA's CUDA-matched cuRobo source package separately when collision-aware planning against an explicit cuRobo world is required. -Use NeuralPlanner (experimental) when you have a trained APG checkpoint and need -learned EEF waypoint rollout on Franka Panda. +Use NeuralPlanner (experimental) when you have a standalone NMG ONNX policy and +need learned EEF waypoint rollout on Franka Panda. See also -------- diff --git a/docs/source/overview/sim/planners/neural_planner.md b/docs/source/overview/sim/planners/neural_planner.md index 632e2c488..03041102f 100644 --- a/docs/source/overview/sim/planners/neural_planner.md +++ b/docs/source/overview/sim/planners/neural_planner.md @@ -3,21 +3,24 @@ ````{admonition} Experimental :class: warning -`NeuralPlanner` is an **experimental** feature. The API, checkpoint format, +`NeuralPlanner` is an **experimental** feature. The API, ONNX model contract, and default parameters may change without a deprecation cycle. It is currently only validated on the **Franka Panda** robot. ```` `NeuralPlanner` is a learning-based EEF waypoint planner. It rolls out a -trained APG checkpoint through `MotionGenerator` to reach Cartesian targets. +standalone NMG ONNX policy through `MotionGenerator` to reach Cartesian targets. +The ONNX graph must include raw-observation normalization. ## Configuration -Pre-trained checkpoints are hosted on HuggingFace and can be downloaded with -`download_neural_planner_checkpoint()` (requires `HF_TOKEN` environment variable). +Install the optional runtime and export the trained policy to ONNX before use: + +```bash +pip install -e '.[nmg]' +``` ```python -from embodichain.data.assets.planner_assets import download_neural_planner_checkpoint from embodichain.lab.sim.planners import ( MotionGenCfg, MotionGenOptions, @@ -28,13 +31,13 @@ from embodichain.lab.sim.planners import ( ) from embodichain.lab.sim.planners.neural_planner import NeuralPlanOptions -checkpoint_path = download_neural_planner_checkpoint() +onnx_model_path = "/path/to/best_mean.onnx" motion_generator = MotionGenerator( cfg=MotionGenCfg( planner_cfg=NeuralPlannerCfg( robot_uid=robot.uid, - checkpoint_path=checkpoint_path, + onnx_model_path=onnx_model_path, control_part="main_arm", ) ) @@ -57,7 +60,18 @@ result = motion_generator.generate( ## Example ```bash -python examples/sim/planners/neural_planner.py --headless --device cuda +python examples/sim/planners/neural_planner.py \ + --headless --device cuda \ + --onnx-model-path /path/to/best_mean.onnx ``` -The example downloads the checkpoint automatically on first run. +The NMG exporter produces a dynamic-batch ONNX policy, so one graph can serve +single-env and env-batched rollout. If the runtime robot base frame or TCP differs +from training, configure `policy_frame_from_world` and +`runtime_tcp_from_policy_tcp` as explicit homogeneous transforms. The conversion is + +```text +policy_T_policy_tcp = policy_T_world + @ world_T_runtime_tcp + @ runtime_tcp_T_policy_tcp +``` diff --git a/embodichain/lab/sim/atomic_actions/core.py b/embodichain/lab/sim/atomic_actions/core.py index 404781920..d487e206d 100644 --- a/embodichain/lab/sim/atomic_actions/core.py +++ b/embodichain/lab/sim/atomic_actions/core.py @@ -545,7 +545,16 @@ def build_plan( raise ValueError("Trajectory and planning context batch sizes must match.") if timed.robot_dof != context.robot.robot_dof: raise ValueError("Trajectory robot_dof must match the planning context.") - timed = timed.hold_rows(success_mask, context.robot.qpos) + planner = ( + None + if self._planning_services is None + else self._planning_services.motion_generator.planner + ) + preserve_failed_positions = ( + getattr(planner, "preserve_failed_plan_positions", False) is True + ) + if not preserve_failed_positions: + timed = timed.hold_rows(success_mask, context.robot.qpos) commands = self._joint_command_sequence( request, diff --git a/embodichain/lab/sim/atomic_actions/engine.py b/embodichain/lab/sim/atomic_actions/engine.py index d86539e79..5c166e057 100644 --- a/embodichain/lab/sim/atomic_actions/engine.py +++ b/embodichain/lab/sim/atomic_actions/engine.py @@ -491,9 +491,18 @@ def compile( f"Skill {plan.skill_id!r} emits non-joint runtime commands and " "cannot be used with offline joint-trajectory compilation." ) - trajectory = plan.joint_trajectory.hold_rows( - step_success, - previous_qpos, + preserve_failed_positions = ( + getattr( + self.motion_generator.planner, + "preserve_failed_plan_positions", + False, + ) + is True + ) + trajectory = ( + plan.joint_trajectory + if preserve_failed_positions + else plan.joint_trajectory.hold_rows(step_success, previous_qpos) ) plans.append(plan) trajectories.append(trajectory) diff --git a/embodichain/lab/sim/planners/base_planner.py b/embodichain/lab/sim/planners/base_planner.py index 5dabfc930..ab09c57ef 100644 --- a/embodichain/lab/sim/planners/base_planner.py +++ b/embodichain/lab/sim/planners/base_planner.py @@ -224,6 +224,9 @@ def __init__(self, cfg: BasePlannerCfg): waypoints for a joint-only backend. """ + supports_heterogeneous_waypoints: bool = False + """Whether one plan may contain an ordered mixture of movement types.""" + preserve_plan_samples: bool = False """Whether callers must retain this planner's returned sample points exactly. diff --git a/embodichain/lab/sim/planners/motion_generator.py b/embodichain/lab/sim/planners/motion_generator.py index a15c224b8..4f3ab626e 100644 --- a/embodichain/lab/sim/planners/motion_generator.py +++ b/embodichain/lab/sim/planners/motion_generator.py @@ -212,6 +212,11 @@ def supports_joint_trajectory_validation(self) -> bool: is True ) + @property + def supports_heterogeneous_waypoints(self) -> bool: + """Whether the selected backend accepts mixed waypoint movement types.""" + return getattr(self.planner, "supports_heterogeneous_waypoints", False) is True + @property def dynamic_collision_entity_ids(self) -> tuple[str, ...]: """Return canonical dynamic-obstacle IDs declared by the planner.""" @@ -488,11 +493,19 @@ def generate( ) move_types = {state.move_type for state in target_states} - if len(move_types) != 1: + heterogeneous = len(move_types) > 1 + if heterogeneous and not self.supports_heterogeneous_waypoints: names = sorted(move_type.name for move_type in move_types) - raise ValueError(f"All target states must share move_type; got {names}.") + raise ValueError( + f"{type(self.planner).__name__} does not support heterogeneous " + f"waypoints; got {names}." + ) + if heterogeneous and options.strategy == "ik_interp": + raise ValueError( + "strategy='ik_interp' does not support heterogeneous waypoints." + ) move_type = target_states[0].move_type - use_interpolation = ( + use_interpolation = not heterogeneous and ( options.preserve_cartesian_samples or options.strategy == "ik_interp" or ( @@ -512,9 +525,11 @@ def _generate_with_planner( options: MotionGenOptions, ) -> PlanResult: """Dispatch batched targets through the configured planner backend.""" + move_types = {state.move_type for state in target_states} move_type = target_states[0].move_type should_preinterpolate = ( - options.is_interpolate + len(move_types) == 1 + and options.is_interpolate and not self.planner.supports_move_type(MoveType.EEF_MOVE) and self.planner.supports_move_type(MoveType.JOINT_MOVE) ) @@ -566,9 +581,11 @@ def _generate_with_planner( else: target_plan_states = target_states - unsupported_move_types = ( - set() if self.planner.supports_move_type(move_type) else {move_type} - ) + unsupported_move_types = { + candidate + for candidate in move_types + if not self.planner.supports_move_type(candidate) + } if not should_preinterpolate and unsupported_move_types: unsupported_names = sorted( move_type.name for move_type in unsupported_move_types @@ -868,7 +885,14 @@ def normalize_derivative( velocities = normalize_derivative(result.velocities, "velocities") accelerations = normalize_derivative(result.accelerations, "accelerations") - if start_qpos is not None and not success.all(): + preserve_failed_positions = ( + getattr(self.planner, "preserve_failed_plan_positions", False) is True + ) + if ( + start_qpos is not None + and not success.all() + and not preserve_failed_positions + ): held = ( start_qpos.to(dtype=positions.dtype).unsqueeze(1).expand_as(positions) ) diff --git a/embodichain/lab/sim/planners/neural_planner.py b/embodichain/lab/sim/planners/neural_planner.py index f5c3d8c90..ddae08740 100644 --- a/embodichain/lab/sim/planners/neural_planner.py +++ b/embodichain/lab/sim/planners/neural_planner.py @@ -17,10 +17,8 @@ from dataclasses import MISSING from pathlib import Path - import numpy as np import torch -import torch.nn as nn from embodichain.lab.sim.planners.base_planner import ( BasePlanner, @@ -40,169 +38,104 @@ ] -def _safe_torch_load(path: Path, map_location: torch.device) -> dict: - """Load a PyTorch checkpoint with safe deserialization when possible. - - Attempts ``weights_only=True`` first. If that fails (e.g. on older PyTorch - versions or checkpoints with unsupported pickle objects), falls back to - ``weights_only=False`` and logs a warning. +class _OnnxPolicy: + """Small ONNX Runtime wrapper that keeps NMG model code out of EmbodiChain.""" - Args: - path: Path to the checkpoint file. - map_location: Device to map tensors to. + def __init__(self, path: Path, providers: list[str] | None = None) -> None: + try: + import onnxruntime as ort + except ImportError as exc: + raise ImportError( + "NeuralPlanner requires onnxruntime. Install EmbodiChain with " + "the 'nmg' optional dependency." + ) from exc - Returns: - The loaded checkpoint dictionary. - """ - try: - return torch.load(path, map_location=map_location, weights_only=True) - except (TypeError, RuntimeError, AttributeError) as exc: - logger.log_warning( - f"Failed to load checkpoint with weights_only=True from {path}: {exc}. " - "Falling back to weights_only=False." + options = ort.SessionOptions() + options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL + self.session = ort.InferenceSession( + str(path), sess_options=options, providers=providers + ) + inputs = self.session.get_inputs() + outputs = self.session.get_outputs() + if len(inputs) != 1 or len(outputs) != 1: + raise ValueError( + "NMG ONNX policy must expose exactly one input and one output; " + f"got {len(inputs)} inputs and {len(outputs)} outputs." + ) + self.input_name = inputs[0].name + self.output_name = outputs[0].name + input_shape = inputs[0].shape + output_shape = outputs[0].shape + if len(input_shape) != 2 or not isinstance(input_shape[1], int): + raise ValueError( + f"Expected ONNX input shape [batch, obs], got {input_shape}." + ) + if len(output_shape) != 2 or output_shape[1] != 7: + raise ValueError( + f"Expected ONNX output shape [batch, 7], got {output_shape}." + ) + self.obs_dim = int(input_shape[1]) + self.fixed_batch_size = ( + int(input_shape[0]) if isinstance(input_shape[0], int) else None ) - return torch.load(path, map_location=map_location, weights_only=False) - - -def _layer_init(layer: nn.Linear, std: float = np.sqrt(2), bias_const: float = 0.0): - torch.nn.init.orthogonal_(layer.weight, std) - torch.nn.init.constant_(layer.bias, bias_const) - return layer - -class _RunningObsNormalizer: - def __init__(self, mean: torch.Tensor, var: torch.Tensor): - self.mean = mean - self.var = var + def __call__(self, obs: torch.Tensor) -> torch.Tensor: + if self.fixed_batch_size is not None and obs.shape[0] != self.fixed_batch_size: + raise ValueError( + f"ONNX policy has fixed batch size {self.fixed_batch_size}, " + f"but received {obs.shape[0]}." + ) + obs_np = np.ascontiguousarray(obs.detach().cpu().numpy(), dtype=np.float32) + action = self.session.run([self.output_name], {self.input_name: obs_np})[0] + return torch.as_tensor(action, dtype=torch.float32, device=obs.device) - def normalize(self, obs: torch.Tensor) -> torch.Tensor: - return (obs - self.mean) / (self.var.sqrt() + 1e-8) +def _waypoint_obs_dim(num_waypoints: int, use_relative_obs: bool) -> int: + """Return the unified NMG constraint-observation width.""" + n = int(num_waypoints) + dim = 7 + 7 + n * (3 + 4 + 7 + 5) + 7 + n + if use_relative_obs: + dim += 7 + n * (3 + 4 + 7) + return dim -class _WaypointTransformerActor(nn.Module): - """APG waypoint actor runtime copied in lightweight form for inference.""" - def __init__( - self, - obs_dim: int, - action_dim: int, - num_waypoints: int, - use_relative_obs: bool = True, - hidden_dim: int = 256, - transformer_nhead: int = 4, - transformer_num_layers: int = 2, - transformer_ff_dim: int | None = None, - ): - super().__init__() - self.num_waypoints = int(num_waypoints) - self.use_relative_obs = bool(use_relative_obs) - if int(action_dim) != 7: - raise ValueError( - "Waypoint transformer checkpoints currently assume a 7-DoF arm. " - f"Got action_dim={action_dim}." - ) - self.state_dim = 7 + 7 + 7 + (7 if self.use_relative_obs else 0) - self.waypoint_token_dim = 3 + 4 + 1 + 1 - - expected_obs_dim = 7 + 7 + self.num_waypoints * 3 - expected_obs_dim += self.num_waypoints * 4 + self.num_waypoints * 2 + 7 - if self.use_relative_obs: - expected_obs_dim += 7 - if int(obs_dim) != expected_obs_dim: - raise ValueError( - "Waypoint transformer expected obs_dim " - f"{expected_obs_dim}, got {obs_dim}." - ) +def _quat_inverse_xyzw(q: torch.Tensor) -> torch.Tensor: + """Return the inverse of a unit quaternion stored as xyzw.""" + return torch.cat([-q[..., :3], q[..., 3:4]], dim=-1) - if hidden_dim % transformer_nhead != 0: - raise ValueError("hidden_dim must be divisible by transformer_nhead") - ff_dim = transformer_ff_dim or hidden_dim * 4 - self.state_proj = _layer_init(nn.Linear(self.state_dim, hidden_dim)) - self.waypoint_proj = _layer_init(nn.Linear(self.waypoint_token_dim, hidden_dim)) - self.state_type_embedding = nn.Parameter(torch.zeros(1, 1, hidden_dim)) - self.waypoint_type_embedding = nn.Parameter(torch.zeros(1, 1, hidden_dim)) - self.waypoint_index_embedding = nn.Parameter( - torch.zeros(1, self.num_waypoints, hidden_dim) - ) +def _quat_mul_xyzw(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + """Multiply quaternion tensors stored as xyzw.""" + ax, ay, az, aw = a.unbind(dim=-1) + bx, by, bz, bw = b.unbind(dim=-1) + return torch.stack( + [ + aw * bx + ax * bw + ay * bz - az * by, + aw * by - ax * bz + ay * bw + az * bx, + aw * bz + ax * by - ay * bx + az * bw, + aw * bw - ax * bx - ay * by - az * bz, + ], + dim=-1, + ) - encoder_layer = nn.TransformerEncoderLayer( - d_model=hidden_dim, - nhead=transformer_nhead, - dim_feedforward=ff_dim, - dropout=0.0, - activation="gelu", - batch_first=True, - norm_first=False, - ) - self.encoder = nn.TransformerEncoder( - encoder_layer, num_layers=transformer_num_layers - ) - self.actor_head = nn.Sequential( - nn.LayerNorm(hidden_dim), - _layer_init(nn.Linear(hidden_dim, hidden_dim)), - nn.Tanh(), - _layer_init(nn.Linear(hidden_dim, action_dim), std=0.01), - ) - def _parse_obs(self, x: torch.Tensor): - n = self.num_waypoints - cursor = 0 - joint = x[:, cursor : cursor + 7] - cursor += 7 - eef_pose = x[:, cursor : cursor + 7] - cursor += 7 - waypoint_pos = x[:, cursor : cursor + 3 * n].reshape(-1, n, 3) - cursor += 3 * n - waypoint_quat = x[:, cursor : cursor + 4 * n].reshape(-1, n, 4) - cursor += 4 * n - active_onehot = x[:, cursor : cursor + n].reshape(-1, n, 1) - cursor += n - valid_mask = x[:, cursor : cursor + n].reshape(-1, n, 1) - cursor += n - last_action = x[:, cursor : cursor + 7] - cursor += 7 - - state_parts = [joint, eef_pose, last_action] - if self.use_relative_obs: - state_parts.append(x[:, cursor : cursor + 7]) - - state = torch.cat(state_parts, dim=-1) - waypoint_tokens = torch.cat( - [waypoint_pos, waypoint_quat, active_onehot, valid_mask], dim=-1 - ) - return state, waypoint_tokens, valid_mask.squeeze(-1) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - state, waypoint_tokens, valid_mask = self._parse_obs(x) - state_token = self.state_proj(state).unsqueeze(1) + self.state_type_embedding - waypoint_tokens = ( - self.waypoint_proj(waypoint_tokens) - + self.waypoint_type_embedding - + self.waypoint_index_embedding - ) - tokens = torch.cat([state_token, waypoint_tokens], dim=1) - state_padding = torch.zeros( - valid_mask.shape[0], 1, dtype=torch.bool, device=valid_mask.device - ) - waypoint_padding = valid_mask < 0.5 - padding_mask = torch.cat([state_padding, waypoint_padding], dim=1) - encoded = self.encoder(tokens, src_key_padding_mask=padding_mask) - return self.actor_head(encoded[:, 0]) +def _canonicalize_quat_xyzw(q: torch.Tensor) -> torch.Tensor: + """Map equivalent quaternion signs to the training-time ``w >= 0`` half.""" + return torch.where(q[..., 3:4] < 0.0, -q, q) @configclass class NeuralPlannerCfg(BasePlannerCfg): planner_type: str = "neural" - checkpoint_path: str = MISSING - """Path to an APG waypoint checkpoint (.pt), e.g. from ``download_neural_planner_checkpoint()``.""" + onnx_model_path: str = MISSING + """Path to a standalone NMG ONNX policy, including observation normalization.""" control_part: str | None = None """Robot control part used for FK and qpos, e.g. 'left_arm'.""" - max_steps: int | None = None - """Maximum rollout steps. If None, uses checkpoint max_episode_steps.""" + max_steps: int = 240 + """Maximum closed-loop kinematic rollout steps.""" action_scale: float = 0.2 """Delta joint scaling factor in radians.""" @@ -210,11 +143,40 @@ class NeuralPlannerCfg(BasePlannerCfg): num_arm_joints: int = 7 """Number of arm joints controlled by the APG policy.""" - pos_eps: float | None = None - """Waypoint position threshold. If None, uses checkpoint waypoint_pos_threshold.""" + num_waypoints: int = 8 + """Number of constraint slots encoded by the ONNX policy.""" + + use_relative_obs: bool = True + """Whether the exported policy uses the unified relative-observation blocks.""" + + canonicalize_quat_obs: bool = True + """Whether quaternion observations use the training-time ``w >= 0`` convention.""" + + intermediate_orientation: bool = True + """Whether every Cartesian waypoint requires its orientation constraint.""" - rot_eps: float | None = None - """Waypoint rotation threshold. If None, uses checkpoint waypoint_rot_threshold.""" + pos_eps: float = 0.01 + """Waypoint position threshold in metres.""" + + rot_eps: float = 0.1 + """Waypoint rotation threshold in radians.""" + + joint_eps: float = 0.02 + """Waypoint joint-position threshold in radians.""" + + onnx_providers: list[str] | None = None + """Optional ONNX Runtime execution-provider priority list.""" + + policy_frame_from_world: list[list[float]] | None = None + """Optional left transform ``policy_T_world`` for runtime poses.""" + + runtime_tcp_from_policy_tcp: list[list[float]] | None = None + """Optional right transform ``runtime_tcp_T_policy_tcp`` for runtime poses. + + Together, pose conversion is + ``policy_T_policy_tcp = policy_T_world @ world_T_runtime_tcp + @ runtime_tcp_T_policy_tcp``. + """ dt: float = 0.01 """Nominal timestep reported in PlanResult.""" @@ -230,9 +192,9 @@ class NeuralPlanOptions(PlanOptions): class NeuralPlanner(BasePlanner): r"""Neural motion planner based on an APG waypoint transformer policy. - The planner loads a checkpoint containing a waypoint-conditioned actor and - rolls it out in closed loop to drive the arm toward a sequence of - end-effector waypoints. Velocities and accelerations in the returned + The planner loads a standalone ONNX waypoint policy and rolls it out in + closed loop to drive the arm toward a sequence of end-effector and/or + joint-position waypoints. Velocities and accelerations in the returned :class:`PlanResult` are estimated via finite differences over the generated position trajectory. @@ -240,20 +202,25 @@ class NeuralPlanner(BasePlanner): cfg: Configuration for the neural planner. Raises: - ValueError: If ``checkpoint_path`` is missing or invalid. - FileNotFoundError: If the checkpoint file does not exist. - KeyError: If the checkpoint is missing required keys. + ValueError: If ``onnx_model_path`` is missing or invalid. + FileNotFoundError: If the ONNX model file does not exist. + ImportError: If ONNX Runtime is not installed. """ - supported_move_types = frozenset({MoveType.EEF_MOVE}) + supported_move_types = frozenset({MoveType.EEF_MOVE, MoveType.JOINT_MOVE}) + supports_heterogeneous_waypoints = True + preserve_plan_samples = True + """Keep native closed-loop states instead of distance-resampling them.""" + preserve_failed_plan_positions = True + """Keep closed-loop rollout samples even when not all waypoints converge.""" def __init__(self, cfg: NeuralPlannerCfg): super().__init__(cfg) self.cfg: NeuralPlannerCfg = cfg - if cfg.checkpoint_path is MISSING or not str(cfg.checkpoint_path): - logger.log_error("checkpoint_path is required", ValueError) - self._load_checkpoint(Path(cfg.checkpoint_path)) + if cfg.onnx_model_path is MISSING or not str(cfg.onnx_model_path): + logger.log_error("onnx_model_path is required", ValueError) + self._load_onnx_model(Path(cfg.onnx_model_path)) def default_plan_options(self) -> NeuralPlanOptions: return NeuralPlanOptions() @@ -274,94 +241,62 @@ def with_motion_context( options.start_qpos = start_qpos return options - def _load_checkpoint(self, checkpoint_path: Path) -> None: - if not checkpoint_path.exists(): - logger.log_error( - f"Checkpoint not found: {checkpoint_path}", FileNotFoundError - ) - - ckpt = _safe_torch_load(checkpoint_path, map_location=self.device) - if "agent" not in ckpt: - raise KeyError( - f"Checkpoint at '{checkpoint_path}' is missing 'agent'. " - f"Available keys: {list(ckpt.keys())}." - ) - if "obs_normalizer" not in ckpt: - raise KeyError( - f"Checkpoint at '{checkpoint_path}' is missing 'obs_normalizer'. " - f"Available keys: {list(ckpt.keys())}." + def _load_onnx_model(self, model_path: Path) -> None: + if not model_path.exists(): + logger.log_error(f"ONNX policy not found: {model_path}", FileNotFoundError) + if model_path.suffix.lower() != ".onnx": + raise ValueError( + "NeuralPlanner only accepts standalone .onnx policies; " + f"got {model_path}." ) - for subkey in ("mean", "var"): - if subkey not in ckpt["obs_normalizer"]: - raise KeyError( - f"Checkpoint obs_normalizer is missing '{subkey}'. " - f"Available: {list(ckpt['obs_normalizer'].keys())}." - ) - self._ckpt_args = ckpt.get("args", {}) - self._num_waypoints = int(self._ckpt_args.get("waypoint_max", 1)) - self._use_relative_obs = bool( - self._ckpt_args.get("waypoint_use_relative_obs", True) - ) - self._policy_arch = self._ckpt_args.get("policy_arch") - if self._policy_arch != "transformer": + self._num_waypoints = int(self.cfg.num_waypoints) + self._use_relative_obs = bool(self.cfg.use_relative_obs) + self._canonicalize_quat_obs = bool(self.cfg.canonicalize_quat_obs) + self._action_dim = int(self.cfg.num_arm_joints) + if self._action_dim != 7: raise ValueError( - "NeuralPlanner only supports transformer waypoint checkpoints. " - f"Got policy_arch={self._policy_arch!r}." + f"NMG ONNX policy controls 7 arm joints, got {self._action_dim}." ) - self._hidden_dim = int(self._ckpt_args.get("hidden_dim", 256)) - self._action_dim = int(self.cfg.num_arm_joints) - self._obs_dim = int(ckpt["obs_normalizer"]["mean"].numel()) - self._max_steps = int( - self.cfg.max_steps or self._ckpt_args.get("max_episode_steps", 30) + self._max_steps = int(self.cfg.max_steps) + self._pos_eps = float(self.cfg.pos_eps) + self._rot_eps = float(self.cfg.rot_eps) + self._joint_eps = float(self.cfg.joint_eps) + self._intermediate_orientation = bool(self.cfg.intermediate_orientation) + self._policy = _OnnxPolicy(model_path, self.cfg.onnx_providers) + self._obs_dim = self._policy.obs_dim + self._policy_frame_from_world = self._as_transform( + self.cfg.policy_frame_from_world, "policy_frame_from_world" ) - self._pos_eps = float( - self.cfg.pos_eps - if self.cfg.pos_eps is not None - else self._ckpt_args.get("waypoint_pos_threshold", 0.05) + self._runtime_tcp_from_policy_tcp = self._as_transform( + self.cfg.runtime_tcp_from_policy_tcp, "runtime_tcp_from_policy_tcp" ) - self._rot_eps = float( - self.cfg.rot_eps - if self.cfg.rot_eps is not None - else self._ckpt_args.get("waypoint_rot_threshold", 0.3) - ) - self._intermediate_orientation = bool( - self._ckpt_args.get("waypoint_intermediate_orientation", True) + expected_obs_dim = _waypoint_obs_dim( + self._num_waypoints, self._use_relative_obs ) + if self._obs_dim != expected_obs_dim: + raise ValueError( + f"ONNX input has obs dim {self._obs_dim}, but the configured " + f"unified constraint layout requires {expected_obs_dim}." + ) - self._normalizer = _RunningObsNormalizer( - ckpt["obs_normalizer"]["mean"].to(self.device), - ckpt["obs_normalizer"]["var"].to(self.device), - ) - self._actor = self._build_actor().to(self.device) - - state_dict = { - k.replace("actor_mean.", ""): v.to(self.device) - for k, v in ckpt["agent"].items() - if k.startswith("actor_mean.") - } - if not state_dict: - raise KeyError("Checkpoint agent has no actor_mean.* weights.") - self._actor.load_state_dict(state_dict) - self._actor.eval() - - def _build_actor(self) -> nn.Module: - return _WaypointTransformerActor( - obs_dim=self._obs_dim, - action_dim=self._action_dim, - num_waypoints=self._num_waypoints, - use_relative_obs=self._use_relative_obs, - hidden_dim=self._hidden_dim, - transformer_nhead=int(self._ckpt_args.get("transformer_nhead", 4)), - transformer_num_layers=int( - self._ckpt_args.get("transformer_num_layers", 2) - ), - transformer_ff_dim=( - int(self._ckpt_args.get("transformer_ff_dim", 0)) or None - ), - ) + def _as_transform(self, value: list[list[float]] | None, name: str) -> torch.Tensor: + """Convert an optional homogeneous-transform config to a device tensor.""" + if value is None: + return torch.eye(4, dtype=torch.float32, device=self.device) + transform = torch.as_tensor(value, dtype=torch.float32, device=self.device) + if transform.shape != (4, 4): + raise ValueError(f"{name} must have shape (4, 4), got {transform.shape}.") + return transform + + def _to_policy_frame(self, xpos: torch.Tensor) -> torch.Tensor: + """Map runtime-world TCP poses to the NMG training base and TCP frame.""" + left = self._policy_frame_from_world.expand(xpos.shape[0], -1, -1) + right = self._runtime_tcp_from_policy_tcp.expand(xpos.shape[0], -1, -1) + return torch.bmm(torch.bmm(left, xpos), right) @validate_plan_options(options_cls=NeuralPlanOptions) + @torch.no_grad() def plan( self, target_states: list[PlanState], @@ -373,9 +308,9 @@ def plan( until all waypoints are reached or ``max_steps`` is exhausted. Args: - target_states: List of :class:`PlanState` waypoints. Each entry must - use :attr:`MoveType.EEF_MOVE` and carry an ``xpos`` tensor of - shape ``(B, 4, 4)``. + target_states: List of :class:`PlanState` waypoints. Each entry uses + :attr:`MoveType.EEF_MOVE` with ``xpos`` shape ``(B, 4, 4)`` or + :attr:`MoveType.JOINT_MOVE` with ``qpos`` shape ``(B, 7)``. options: :class:`NeuralPlanOptions` with ``control_part``, ``start_qpos``, and ``max_steps`` overrides. @@ -388,8 +323,8 @@ def plan( via finite differences and are therefore approximate. Raises: - ValueError: If ``control_part`` is not provided, if a target state - is not ``EEF_MOVE``, or if ``start_qpos`` has too few joints. + ValueError: If ``control_part`` is not provided, a target state is + unsupported, or ``start_qpos`` has too few joints. """ if not target_states: return PlanResult( @@ -404,9 +339,16 @@ def plan( ValueError, ) - waypoints_pos, waypoints_quat, valid_mask, episode_k = self._parse_waypoints( - target_states - ) + ( + waypoints_pos, + waypoints_quat, + waypoints_joint, + valid_mask, + pos_mask, + rot_mask, + joint_mask, + episode_k, + ) = self._parse_waypoints(target_states) qpos = self._initial_qpos(control_part, options.start_qpos) b = qpos.shape[0] limits = self.robot.get_qpos_limits(name=control_part)[0].to(self.device) @@ -428,11 +370,15 @@ def plan( ee_pose, waypoints_pos, waypoints_quat, + waypoints_joint, valid_mask, + pos_mask, + rot_mask, + joint_mask, active_idx, last_action, ) - action = self._actor(self._normalizer.normalize(obs)).clamp(-1.0, 1.0) + action = self._policy(obs).clamp(-1.0, 1.0) # Hold converged envs: zero their action so qpos does not drift. # `converged` reflects state up to the end of the previous step, so # once an env converged at the end of step N its action is masked @@ -453,7 +399,15 @@ def plan( ee_pose = self._fk_pose_xyzw(qpos, control_part) reached = self._is_active_reached( - ee_pose, waypoints_pos, waypoints_quat, active_idx, episode_k + qpos[:, : self._action_dim], + ee_pose, + waypoints_pos, + waypoints_quat, + waypoints_joint, + pos_mask, + rot_mask, + joint_mask, + active_idx, ) active_idx = torch.where(reached, active_idx + 1, active_idx) converged = converged | (active_idx >= episode_k) @@ -485,34 +439,87 @@ def plan( dt=dt, ) - def _parse_waypoints( - self, target_states: list[PlanState] - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: + def _parse_waypoints(self, target_states: list[PlanState]) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + int, + ]: if len(target_states) > self._num_waypoints: logger.log_error( - f"Received {len(target_states)} waypoints, but checkpoint supports " + f"Received {len(target_states)} waypoints, but the ONNX policy supports " f"at most {self._num_waypoints}.", ValueError, ) b = _infer_batch_size(target_states) or 1 waypoint_pos = torch.zeros(b, self._num_waypoints, 3, device=self.device) waypoint_quat = torch.zeros(b, self._num_waypoints, 4, device=self.device) + waypoint_quat[..., 3] = 1.0 + waypoint_joint = torch.zeros( + b, self._num_waypoints, self._action_dim, device=self.device + ) valid_mask = torch.zeros(b, self._num_waypoints, device=self.device) + pos_mask = torch.zeros_like(valid_mask) + rot_mask = torch.zeros_like(valid_mask) + joint_mask = torch.zeros_like(valid_mask) for idx, target in enumerate(target_states): - if target.move_type != MoveType.EEF_MOVE or target.xpos is None: + if target.move_type == MoveType.EEF_MOVE and target.xpos is not None: + xpos = torch.as_tensor( + target.xpos, dtype=torch.float32, device=self.device + ) + if xpos.dim() == 2: + xpos = xpos.unsqueeze(0) + policy_xpos = self._to_policy_frame(xpos) + waypoint_pos[:, idx] = policy_xpos[:, :3, 3] + quat_xyzw = convert_quat( + quat_from_matrix(policy_xpos[:, :3, :3]), to="xyzw" + ) + waypoint_quat[:, idx] = ( + _canonicalize_quat_xyzw(quat_xyzw) + if getattr(self, "_canonicalize_quat_obs", False) + else quat_xyzw + ) + pos_mask[:, idx] = 1.0 + rot_mask[:, idx] = 1.0 + elif target.move_type == MoveType.JOINT_MOVE and target.qpos is not None: + qpos = torch.as_tensor( + target.qpos, dtype=torch.float32, device=self.device + ) + if qpos.dim() == 1: + qpos = qpos.unsqueeze(0) + if qpos.shape != (b, self._action_dim): + logger.log_error( + "NeuralPlanner JOINT_MOVE qpos must have shape " + f"({b}, {self._action_dim}), got {tuple(qpos.shape)}.", + ValueError, + ) + waypoint_joint[:, idx] = qpos + joint_mask[:, idx] = 1.0 + else: logger.log_error( - "NeuralPlanner expects EEF_MOVE PlanState entries with xpos.", + "NeuralPlanner expects EEF_MOVE entries with xpos or " + "JOINT_MOVE entries with qpos.", ValueError, ) - xpos = torch.as_tensor(target.xpos, dtype=torch.float32, device=self.device) - if xpos.dim() == 2: - xpos = xpos.unsqueeze(0) - waypoint_pos[:, idx] = xpos[:, :3, 3] - waypoint_quat[:, idx] = convert_quat( - quat_from_matrix(xpos[:, :3, :3]), to="xyzw" - ) valid_mask[:, idx] = 1.0 - return waypoint_pos, waypoint_quat, valid_mask, len(target_states) + if not self._intermediate_orientation: + final_mask = torch.zeros_like(rot_mask) + final_mask[:, len(target_states) - 1] = 1.0 + rot_mask *= final_mask + return ( + waypoint_pos, + waypoint_quat, + waypoint_joint, + valid_mask, + pos_mask, + rot_mask, + joint_mask, + len(target_states), + ) def _initial_qpos( self, control_part: str, start_qpos: torch.Tensor | None @@ -535,9 +542,11 @@ def _fk_matrix(self, qpos: torch.Tensor, control_part: str) -> torch.Tensor: return self.robot.compute_fk(qpos=qpos, name=control_part, to_matrix=True) def _fk_pose_xyzw(self, qpos: torch.Tensor, control_part: str) -> torch.Tensor: - fk = self.robot.compute_fk(qpos=qpos, name=control_part, to_matrix=False) - pos = fk[:, :3] - quat_xyzw = convert_quat(fk[:, 3:7], to="xyzw") + fk = self._to_policy_frame(self._fk_matrix(qpos, control_part)) + pos = fk[:, :3, 3] + quat_xyzw = convert_quat(quat_from_matrix(fk[:, :3, :3]), to="xyzw") + if getattr(self, "_canonicalize_quat_obs", False): + quat_xyzw = _canonicalize_quat_xyzw(quat_xyzw) return torch.cat([pos, quat_xyzw], dim=-1) def _build_obs( @@ -546,7 +555,11 @@ def _build_obs( ee_pose: torch.Tensor, waypoint_pos: torch.Tensor, waypoint_quat: torch.Tensor, + waypoint_joint: torch.Tensor, valid_mask: torch.Tensor, + pos_mask: torch.Tensor, + rot_mask: torch.Tensor, + joint_mask: torch.Tensor, active_idx: torch.Tensor, last_action: torch.Tensor, ) -> torch.Tensor: @@ -554,22 +567,97 @@ def _build_obs( active_idx_clamped = torch.clamp(active_idx, max=self._num_waypoints - 1) active_onehot = torch.zeros(b, self._num_waypoints, device=self.device) active_onehot.scatter_(1, active_idx_clamped.unsqueeze(1), 1.0) + pos_block = waypoint_pos * pos_mask.unsqueeze(-1) + joint_block = waypoint_joint * joint_mask.unsqueeze(-1) + identity = torch.tensor( + [0.0, 0.0, 0.0, 1.0], + dtype=waypoint_quat.dtype, + device=self.device, + ) + quat_block = torch.where( + rot_mask.unsqueeze(-1) > 0.5, + waypoint_quat, + identity.view(1, 1, 4), + ) obs_parts = [ joint_pos, ee_pose, - waypoint_pos.reshape(b, self._num_waypoints * 3), - waypoint_quat.reshape(b, self._num_waypoints * 4), + pos_block.reshape(b, self._num_waypoints * 3), + quat_block.reshape(b, self._num_waypoints * 4), + joint_block.reshape(b, self._num_waypoints * self._action_dim), active_onehot, valid_mask, + pos_mask, + rot_mask, + joint_mask, last_action, ] if self._use_relative_obs: idx = torch.arange(b, device=self.device) - active_pos = waypoint_pos[idx, active_idx_clamped] - active_quat = waypoint_quat[idx, active_idx_clamped] - obs_parts.append( - torch.cat([active_pos - ee_pose[:, :3], active_quat], dim=-1) + active_pos = pos_block[idx, active_idx_clamped] + active_quat = quat_block[idx, active_idx_clamped] + active_joint = joint_block[idx, active_idx_clamped] + inv_eef = _quat_inverse_xyzw(ee_pose[:, 3:7]) + active_rel_quat = _quat_mul_xyzw(active_quat, inv_eef) + if getattr(self, "_canonicalize_quat_obs", False): + active_rel_quat = _canonicalize_quat_xyzw(active_rel_quat) + active_pos_mask = pos_mask[idx, active_idx_clamped].unsqueeze(-1) + active_rot_mask = rot_mask[idx, active_idx_clamped].unsqueeze(-1) + active_rel_quat = torch.where( + active_rot_mask > 0.5, + active_rel_quat, + identity.view(1, 4), + ) + active_cart_rel = torch.cat( + [ + (active_pos - ee_pose[:, :3]) * active_pos_mask, + active_rel_quat, + ], + dim=-1, + ) + active_rel = torch.where( + (joint_mask[idx, active_idx_clamped] > 0.5).unsqueeze(-1), + active_joint - joint_pos, + active_cart_rel, + ) + obs_parts.append(active_rel) + rel_pos = (pos_block - ee_pose[:, None, :3]) * pos_mask.unsqueeze(-1) + rel_quat = _quat_mul_xyzw( + quat_block, + inv_eef[:, None, :].expand_as(quat_block), ) + if getattr(self, "_canonicalize_quat_obs", False): + rel_quat = _canonicalize_quat_xyzw(rel_quat) + rel_quat = torch.where( + rot_mask.unsqueeze(-1) > 0.5, + rel_quat, + identity.view(1, 1, 4), + ) + joint_err = (joint_block - joint_pos[:, None]) * joint_mask.unsqueeze(-1) + obs_parts.extend( + [ + rel_pos.reshape(b, self._num_waypoints * 3), + rel_quat.reshape(b, self._num_waypoints * 4), + joint_err.reshape(b, self._num_waypoints * self._action_dim), + ] + ) + waypoint_type = torch.zeros( + b, + self._num_waypoints, + dtype=joint_pos.dtype, + device=self.device, + ) + waypoint_type = torch.where( + (valid_mask > 0.5) & (rot_mask < 0.5), + torch.ones_like(waypoint_type), + waypoint_type, + ) + waypoint_type = torch.where( + joint_mask > 0.5, + torch.full_like(waypoint_type, 2.0), + waypoint_type, + ) + obs_parts.append(waypoint_type) obs = torch.cat(obs_parts, dim=-1) if obs.shape[-1] != self._obs_dim: raise ValueError( @@ -579,30 +667,46 @@ def _build_obs( def _is_active_reached( self, + joint_pos: torch.Tensor, ee_pose: torch.Tensor, waypoint_pos: torch.Tensor, waypoint_quat: torch.Tensor, + waypoint_joint: torch.Tensor, + pos_mask: torch.Tensor, + rot_mask: torch.Tensor, + joint_mask: torch.Tensor, active_idx: torch.Tensor, - episode_k: int, ) -> torch.Tensor: b = ee_pose.shape[0] idx = torch.arange(b, device=self.device) active_idx_clamped = torch.clamp(active_idx, max=self._num_waypoints - 1) active_pos = waypoint_pos[idx, active_idx_clamped] active_quat_xyzw = waypoint_quat[idx, active_idx_clamped] + active_joint = waypoint_joint[idx, active_idx_clamped] + active_pos_mask = pos_mask[idx, active_idx_clamped] > 0.5 + active_rot_mask = rot_mask[idx, active_idx_clamped] > 0.5 + active_joint_mask = joint_mask[idx, active_idx_clamped] > 0.5 pos_dist = (ee_pose[:, :3] - active_pos).norm(dim=-1) ee_quat_wxyz = convert_quat(ee_pose[:, 3:7], to="wxyz") active_quat_wxyz = convert_quat(active_quat_xyzw, to="wxyz") rot_dist = quat_error_magnitude(ee_quat_wxyz, active_quat_wxyz) - orientation_required = self._intermediate_orientation | ( - active_idx >= episode_k - 1 - ) rot_ok = torch.where( - orientation_required, + active_rot_mask, rot_dist < self._rot_eps, torch.ones_like(rot_dist, dtype=torch.bool), ) - reached = (pos_dist < self._pos_eps) & rot_ok + pos_ok = torch.where( + active_pos_mask, + pos_dist < self._pos_eps, + torch.ones_like(pos_dist, dtype=torch.bool), + ) + joint_dist = torch.amax(torch.abs(joint_pos - active_joint), dim=-1) + joint_ok = torch.where( + active_joint_mask, + joint_dist < self._joint_eps, + torch.ones_like(joint_dist, dtype=torch.bool), + ) + reached = pos_ok & rot_ok & joint_ok return reached @staticmethod diff --git a/examples/sim/planners/neural_planner.py b/examples/sim/planners/neural_planner.py index 9ce5a18db..e326933cd 100644 --- a/examples/sim/planners/neural_planner.py +++ b/examples/sim/planners/neural_planner.py @@ -14,12 +14,12 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Run the env-batched NeuralPlanner waypoint example. +"""Run the ONNX NeuralPlanner waypoint example. From the repository root:: - python examples/sim/planners/neural_planner.py --headless - python examples/sim/planners/neural_planner.py --headless --device cuda:1 + python examples/sim/planners/neural_planner.py --headless \ + --onnx-model-path /path/to/policy.onnx """ from __future__ import annotations @@ -30,7 +30,6 @@ import numpy as np import torch -from embodichain.data.assets.planner_assets import download_neural_planner_checkpoint from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.visualization import visualization_cfg_from_args @@ -53,6 +52,11 @@ def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="NeuralPlanner waypoint example") add_env_launcher_args_to_parser(parser) parser.set_defaults(device=default_device, arena_space=2.0) + parser.add_argument( + "--onnx-model-path", + required=True, + help="Path to a standalone NMG ONNX policy.", + ) parser.add_argument( "--num-waypoints", type=int, @@ -188,15 +192,15 @@ def play_trajectory( def main() -> None: args = parse_args() - if args.num_envs < 1: - raise ValueError("--num_envs must be at least 1.") + if args.num_envs != 1: + raise ValueError("The current exported NMG ONNX policy requires --num-envs 1.") if args.num_waypoints < 1: raise ValueError("--num-waypoints must be at least 1.") if args.step_repeat < 1: raise ValueError("--step-repeat must be at least 1.") if args.hold_steps < 0: raise ValueError("--hold-steps must be non-negative.") - checkpoint_path = download_neural_planner_checkpoint() + onnx_model_path = args.onnx_model_path sim_device = _resolve_device(args.device, args.gpu_id) resolved_device = torch.device(sim_device) @@ -246,7 +250,7 @@ def main() -> None: cfg=MotionGenCfg( planner_cfg=NeuralPlannerCfg( robot_uid=robot.uid, - checkpoint_path=checkpoint_path, + onnx_model_path=onnx_model_path, control_part=arm_name, ) ) diff --git a/pyproject.toml b/pyproject.toml index faece9731..be5ff76aa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,6 +52,9 @@ dependencies = [ ] [project.optional-dependencies] +nmg = [ + "onnxruntime>=1.20", +] gensim = [ "bpy", "gradio>=6.17.3,<6.18", diff --git a/scripts/benchmark/__main__.py b/scripts/benchmark/__main__.py index 226843aa4..68ccc98f4 100644 --- a/scripts/benchmark/__main__.py +++ b/scripts/benchmark/__main__.py @@ -51,7 +51,7 @@ def _run_rl_cli(_: argparse.Namespace) -> None: def _run_motion_generation_cli(args: argparse.Namespace) -> None: - """Run the free-space motion-generation benchmark.""" + """Run the shared planner and Atomic Task benchmark.""" from scripts.benchmark.motion_generation.run_benchmark import run_from_args run_from_args(args) @@ -118,7 +118,7 @@ def main(argv: Sequence[str] | None = None) -> None: motion_generation_parser = subparsers.add_parser( "motion-generation", - help="Benchmark free-space motion generation with cuRobo as baseline.", + help="Benchmark planners on fixed motion and Atomic Task cases.", ) add_parser_arguments(motion_generation_parser) motion_generation_parser.set_defaults(func=_run_motion_generation_cli) diff --git a/scripts/benchmark/motion_generation/BENCHMARK_DESIGN.md b/scripts/benchmark/motion_generation/BENCHMARK_DESIGN.md index ae1a1bddc..d53732eec 100644 --- a/scripts/benchmark/motion_generation/BENCHMARK_DESIGN.md +++ b/scripts/benchmark/motion_generation/BENCHMARK_DESIGN.md @@ -15,6 +15,28 @@ The default comparison should be NMG versus cuRobo. IK plus interpolation and TOPPRA should remain optional diagnostic baselines rather than define the main leaderboard. +## Implemented vertical slice + +The first physics-backed slice is available as +`suites/atomic_franka_pgi_curobo.yaml`. It runs Franka + PGI with cuRobo only, +and covers `MoveEndEffector` plus antipodal-grasp `PickUp`. Both skills pin +`MotionPolicy(strategy="motion_gen", planner="curobo")` and compile through the +same `AtomicActionEngine`; scenario code never calls cuRobo directly. + +The shared runner now selects planners, scenarios, robots, Atomic Action case +providers, and object kinds through registries. Cases freeze the full robot +start state, explicit targets/grasp, object configuration, difficulty factors, +and independent sequential-IK evidence before measured planner calls. Reports +keep planning, kinematic motion validity, controller execution, and physical +task success separate while retaining exactly three tables. + +Run the slice with: + +```bash +python -m scripts.benchmark.motion_generation.run_benchmark \ + --suite atomic_franka_pgi_curobo --device cuda +``` + ## Motivation The existing NeuralPlanner benchmark provides useful latency, memory, rollout, diff --git a/scripts/benchmark/motion_generation/README.md b/scripts/benchmark/motion_generation/README.md index 9ec7ce42d..47671cb82 100644 --- a/scripts/benchmark/motion_generation/README.md +++ b/scripts/benchmark/motion_generation/README.md @@ -1,36 +1,66 @@ -# Motion Generation Benchmark +# Planner Motion Generation & Atomic Skill Benchmark -Free-space motion-generation suite with cuRobo as the default primary baseline. +Shared planner benchmark framework for fixed motion-generation cases and +physics-backed Atomic Actions. All Atomic Actions call the selected planner +through the adapter-owned `MotionGenerator`. Design background and roadmap: see [`BENCHMARK_DESIGN.md`](./BENCHMARK_DESIGN.md). ## Run ```bash -embodichain benchmark motion-generation --suite smoke -embodichain benchmark motion-generation --suite coverage -embodichain benchmark motion-generation --extra-baselines ik_interpolate toppra -embodichain benchmark motion-generation --path-shapes direct l_turn --start-state-bins nominal near_singularity +python -m scripts.benchmark.motion_generation.run_benchmark --suite smoke +python -m scripts.benchmark.motion_generation.run_benchmark --suite coverage +python -m scripts.benchmark.motion_generation.run_benchmark \ + --suite atomic_franka_pgi_curobo --device cuda +python -m scripts.benchmark.motion_generation.run_benchmark \ + --suite atomic_franka_pgi_curobo --device cuda --record-video +python -m scripts.benchmark.motion_generation.run_benchmark \ + --suite atomic_franka_pgi_curobo_randomized --device cuda +python -m scripts.benchmark.motion_generation.run_benchmark \ + --extra-baselines ik_interpolate toppra ``` -Artifacts land under `outputs/benchmarks/motion_generation//` +Artifacts land under `outputs/benchmarks///` (`resolved_suite.yaml`, `case_manifest.json`, `trials.jsonl`, `aggregates.json`, -`report.md` with exactly three tables). +`report.md` with exactly three tables). Atomic Task videos, when enabled, land +in that run's `videos/` directory. ## Implemented -- Extensible planner/scenario registries and track-based suite YAML +- Extensible planner, scenario, robot, Atomic Action, and object registries - `free-space-common` track with fixed manifests and start-state bins +- `atomic-task` track with frozen robot/object/task manifests and common physics replay +- Single-arm Atomic Task slice: Franka + PGI with `MoveEndEffector`, + `MoveJoints`, `PickUp`, `MoveHeldObject`, `Place`, and `Press` +- Deterministic 16-seed generalization sweep for all six single-arm skills, + with bounded robot-start, target, held-object, and object-pose randomization - Default matrix: cuRobo (`primary_baseline`); IK / TOPPRA optional diagnostics -- NMG adapter stub (`candidate`, disabled until a checkpoint is ready) +- Direct, batched NMG ONNX adapter (`candidate`, enabled when a model path is supplied) - Lifecycle timing: construct / prepare / cold / warm -- Ordered waypoint matching and external `motion_valid` (separate from - `PlanResult.success`) +- Distinct planning, motion-valid, execution, and physical task-success stages +- Planning latency, execution wall time, end-to-end time, nominal trajectory + duration, simulated task-completion time, controller tracking RMSE, and + task-specific object lift - One Markdown report: Time & Memory, Success & Other Metrics, Leaderboard +- Optional Atomic Task headless replay videos after measured evaluation + (`--record-video`, `--record-failed-video`) -## Not implemented yet +## Extend -- Real NMG checkpoint adapter -- `collision-deployment` and `atomic-task` tracks -- Physics execution / task-success metrics +- Planner: register a `PlannerAdapter`, expose its `MotionGenerator`, and + declare the `atomic_action` capability. +- Robot: register a `RobotProvider` and select it under `robot`; PickUp suites + also declare the gripper control part and open/grasp qpos under `gripper`. +- Object: add another `objects` entry for built-in `cube`/`mesh`, or register a + new object-kind factory. +- Atomic skill: implement and register an `AtomicSkillCaseProvider`; the runner, + artifact schema, aggregation, and report stay unchanged. + +## Current limits + +- `collision-deployment` and obstacle-aware common-input tracks +- Atomic Task execution is currently `B=1`; the supplied suite covers only + Franka + PGI and cuRobo +- Dual-arm actions and multi-action chains - Latency-budget Pareto sweeps, confidence intervals, subprocess isolation diff --git a/scripts/benchmark/motion_generation/aggregation.py b/scripts/benchmark/motion_generation/aggregation.py index 25dbd45db..60eea4196 100644 --- a/scripts/benchmark/motion_generation/aggregation.py +++ b/scripts/benchmark/motion_generation/aggregation.py @@ -70,6 +70,51 @@ def _case_macro_rate( return sum(case_rates) / len(case_rates) +def _case_supports_attribute(case: BenchmarkCase, attribute: str) -> bool: + """Return whether a staged outcome applies to one case protocol.""" + if attribute == "task_success": + return case.primary_success == "task_success" + if attribute == "execution_success": + return case.primary_success in {"execution_success", "task_success"} + return True + + +def _case_macro_optional_rate( + measured: list[TrialRecord], + track_cases: list[BenchmarkCase], + attribute: str, +) -> float | None: + """Macro-average an applicable staged boolean, otherwise return ``None``.""" + applicable = [ + case for case in track_cases if _case_supports_attribute(case, attribute) + ] + if not applicable: + return None + return _case_macro_rate(measured, applicable, attribute) + + +def _case_macro_primary_rate( + measured: list[TrialRecord], track_cases: list[BenchmarkCase] +) -> float: + """Macro-average each case's explicitly declared primary success stage.""" + if not track_cases: + return 0.0 + outcomes_by_case: dict[str, list[CaseOutcome]] = defaultdict(list) + for record in measured: + outcomes_by_case[record.case_id].extend(record.outcomes) + case_rates: list[float] = [] + for case in track_cases: + outcomes = outcomes_by_case.get(case.case_id, []) + if not outcomes: + case_rates.append(0.0) + continue + case_rates.append( + sum(bool(getattr(outcome, case.primary_success)) for outcome in outcomes) + / len(outcomes) + ) + return sum(case_rates) / len(case_rates) + + def _case_macro_mean( measured: list[TrialRecord], track_cases: list[BenchmarkCase], @@ -183,15 +228,20 @@ def _performance_rows( cases: list[BenchmarkCase], ) -> list[dict[str, object]]: """Aggregate steady-state time and memory by track, algorithm, and input shape.""" - measured_groups: dict[tuple[str, str, int, int], list[TrialRecord]] = defaultdict( - list - ) + measured_groups: dict[ + tuple[str, str, str, str, str | None, str | None, int, int], + list[TrialRecord], + ] = defaultdict(list) for record in records: if record.phase is TrialPhase.MEASURED: measured_groups[ ( record.track, record.algorithm_id, + record.robot_id, + record.skill_id, + record.object_id, + record.task_difficulty, record.batch_size, record.waypoint_count, ) @@ -199,8 +249,29 @@ def _performance_rows( metadata_by_id = {item.algorithm_id: item for item in metadata} rows: list[dict[str, object]] = [] - for key in sorted(measured_groups): - track, algorithm_id, batch_size, waypoint_count = key + for key in sorted( + measured_groups, + key=lambda item: ( + item[0], + item[1], + item[2], + item[3], + item[4] or "", + item[5] or "", + item[6], + item[7], + ), + ): + ( + track, + algorithm_id, + robot_id, + skill_id, + object_id, + task_difficulty, + batch_size, + waypoint_count, + ) = key group = measured_groups[key] info = metadata_by_id[algorithm_id] costs = [record.cost_time_ms for record in group] @@ -210,6 +281,10 @@ def _performance_rows( "track": track, "algorithm": algorithm_id, "algorithm_role": info.algorithm_role.value, + "robot": robot_id, + "skill": skill_id, + "object": object_id, + "task_difficulty": task_difficulty, "batch_size": batch_size, "waypoint_count": waypoint_count, "num_trials": len(group), @@ -244,6 +319,23 @@ def _performance_rows( "cpu_delta_mb": _mean(record.cpu_delta_mb for record in group), "gpu_delta_mb": _mean(record.gpu_delta_mb for record in group), "peak_gpu_mb": _peak_gpu(group), + "execution_time_ms": _mean( + record.execution_time_ms for record in group + ), + "end_to_end_time_ms": _mean( + record.end_to_end_time_ms for record in group + ), + "trajectory_duration_s": _mean( + record.trajectory_duration_s for record in group + ), + "trajectory_waypoints": _mean( + ( + float(record.trajectory_waypoints) + if record.trajectory_waypoints is not None + else None + ) + for record in group + ), } ) @@ -257,6 +349,10 @@ def _performance_rows( "track": track, "algorithm": info.algorithm_id, "algorithm_role": info.algorithm_role.value, + "robot": None, + "skill": None, + "object": None, + "task_difficulty": None, "batch_size": None, "waypoint_count": None, "num_trials": 0, @@ -272,6 +368,10 @@ def _performance_rows( "cpu_delta_mb": None, "gpu_delta_mb": None, "peak_gpu_mb": None, + "execution_time_ms": None, + "end_to_end_time_ms": None, + "trajectory_duration_s": None, + "trajectory_waypoints": None, } ) return sorted( @@ -281,6 +381,8 @@ def _performance_rows( str(row["algorithm"]), int(row["batch_size"] or 0), int(row["waypoint_count"] or 0), + str(row["skill"] or ""), + str(row["object"] or ""), ), ) @@ -298,7 +400,21 @@ def _metric_rows( remain conditioned on externally motion-valid outcomes. """ measured_by_key: dict[ - tuple[str, str, str, int, int, str, str], list[TrialRecord] + tuple[ + str, + str, + str, + str, + str, + str | None, + str | None, + str, + int, + int, + str, + str, + ], + list[TrialRecord], ] = defaultdict(list) for record in records: if record.phase is not TrialPhase.MEASURED: @@ -307,6 +423,11 @@ def _metric_rows( record.track, record.algorithm_id, record.scenario_id, + record.robot_id, + record.skill_id, + record.object_id, + record.task_difficulty, + record.primary_success, record.batch_size, record.waypoint_count, record.path_shape, @@ -314,14 +435,46 @@ def _metric_rows( ) measured_by_key[key].append(record) - expected_by_group: Counter[tuple[str, str, int, int, str, str]] = Counter() - cases_by_group: dict[tuple[str, str, int, int, str, str], list[BenchmarkCase]] = ( - defaultdict(list) - ) + expected_by_group: Counter[ + tuple[ + str, + str, + str, + str, + str | None, + str | None, + str, + int, + int, + str, + str, + ] + ] = Counter() + cases_by_group: dict[ + tuple[ + str, + str, + str, + str, + str | None, + str | None, + str, + int, + int, + str, + str, + ], + list[BenchmarkCase], + ] = defaultdict(list) for case in cases: key = ( case.track, case.scenario_id, + case.robot_id, + case.skill_id, + case.object_id, + case.task_difficulty, + case.primary_success, case.batch_size, case.num_waypoints, case.path_shape, @@ -332,10 +485,30 @@ def _metric_rows( rows: list[dict[str, object]] = [] for info in metadata: - for group_key in sorted(expected_by_group): + for group_key in sorted( + expected_by_group, + key=lambda item: ( + item[0], + item[1], + item[2], + item[3], + item[4] or "", + item[5] or "", + item[6], + item[7], + item[8], + item[9], + item[10], + ), + ): ( track, scenario_id, + robot_id, + skill_id, + object_id, + task_difficulty, + primary_success, batch_size, waypoint_count, path_shape, @@ -347,6 +520,11 @@ def _metric_rows( track, info.algorithm_id, scenario_id, + robot_id, + skill_id, + object_id, + task_difficulty, + primary_success, batch_size, waypoint_count, path_shape, @@ -363,6 +541,11 @@ def _metric_rows( "scenario": scenario_id, "algorithm": info.algorithm_id, "algorithm_role": info.algorithm_role.value, + "robot": robot_id, + "skill": skill_id, + "object": object_id, + "task_difficulty": task_difficulty, + "primary_success": primary_success, "batch_size": batch_size, "waypoint_count": waypoint_count, "path_shape": path_shape, @@ -370,13 +553,19 @@ def _metric_rows( "cases": len(group_cases), "n_valid": len(valid_outcomes), "coverage_rate": min(1.0, len(outcomes) / max(expected, 1)), - # Free-space primary success is external motion validity. - "success_rate": _case_macro_rate( - measured, group_cases, "motion_valid" - ), + "success_rate": _case_macro_primary_rate(measured, group_cases), "planning_success_rate": _case_macro_rate( measured, group_cases, "planning_success" ), + "motion_valid_rate": _case_macro_rate( + measured, group_cases, "motion_valid" + ), + "execution_success_rate": _case_macro_optional_rate( + measured, group_cases, "execution_success" + ), + "task_success_rate": _case_macro_optional_rate( + measured, group_cases, "task_success" + ), "ordered_waypoint_success_rate": _case_macro_rate( measured, group_cases, "ordered_waypoints_reached" ), @@ -409,6 +598,25 @@ def _metric_rows( "path_efficiency": _mean( outcome.path_efficiency for outcome in valid_outcomes ), + "task_completion_time_s": _mean( + outcome.task_completion_time_s + for outcome in outcomes + if outcome.task_success + ), + "joint_tracking_rmse_rad": _mean( + outcome.joint_tracking_rmse_rad for outcome in outcomes + ), + "object_lift_delta_m": _mean( + outcome.object_lift_delta_m for outcome in outcomes + ), + "replan_count": _mean( + ( + float(outcome.replan_count) + if outcome.replan_count is not None + else None + ) + for outcome in outcomes + ), "top_failure": _top_failure(outcomes), } ) @@ -447,6 +655,11 @@ def _leaderboard_rows( coverage = min(1.0, len(outcomes) / max(expected_outcomes, 1)) motion_rate = _case_macro_rate(measured, track_cases, "motion_valid") planning_rate = _case_macro_rate(measured, track_cases, "planning_success") + execution_rate = _case_macro_optional_rate( + measured, track_cases, "execution_success" + ) + task_rate = _case_macro_optional_rate(measured, track_cases, "task_success") + primary_rate = _case_macro_primary_rate(measured, track_cases) latency_p95 = _case_macro_latency_p95(measured, track_cases) peak_gpu = _peak_gpu(measured) track_entries.append( @@ -458,11 +671,11 @@ def _leaderboard_rows( "planner_config_hash": info.config_hash[:12], "eligible": coverage >= 1.0 - 1.0e-12, "coverage_rate": coverage, - # free-space v1: primary_success == motion_valid - "overall_success_rate": motion_rate, + "overall_success_rate": primary_rate, "planning_success_rate": planning_rate, "motion_valid_rate": motion_rate, - "task_success_rate": None, + "execution_success_rate": execution_rate, + "task_success_rate": task_rate, "latency_p95_ms": latency_p95, "peak_gpu_mb": peak_gpu, } diff --git a/scripts/benchmark/motion_generation/artifacts.py b/scripts/benchmark/motion_generation/artifacts.py index 4544230f6..ecda92707 100644 --- a/scripts/benchmark/motion_generation/artifacts.py +++ b/scripts/benchmark/motion_generation/artifacts.py @@ -124,21 +124,51 @@ def _case_to_dict(case: BenchmarkCase) -> dict[str, Any]: "num_waypoints": case.num_waypoints, "path_shape": case.path_shape, "start_state_bin": case.start_state_bin, + "robot_id": case.robot_id, + "skill_id": case.skill_id, + "object_id": case.object_id, + "task_difficulty": case.task_difficulty, + "primary_success": case.primary_success, "start_qpos": case.start_qpos.detach().cpu().tolist(), + "full_start_qpos": ( + None + if case.full_start_qpos is None + else case.full_start_qpos.detach().cpu().tolist() + ), "target_waypoints": case.target_waypoints.detach().cpu().tolist(), + "case_parameters": _to_json_value(case.case_parameters), "validity_evidence": { - "method": "reference_qpos_fk", + "method": ( + "reference_qpos_fk" + if case.skill_id == "N/A" + else ( + "joint_limits_and_reference_fk" + if case.skill_id == "move_joints" + else "independent_sequential_ik" + ) + ), "reference_qpos": case.reference_qpos.detach().cpu().tolist(), }, } +def _to_json_value(value: object) -> object: + """Recursively preserve tensors and numeric case configuration values.""" + if isinstance(value, torch.Tensor): + return value.detach().cpu().tolist() + if isinstance(value, dict): + return {str(key): _to_json_value(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_to_json_value(item) for item in value] + return value + + def write_case_manifest(path: str | Path, cases: list[BenchmarkCase]) -> Path: """Write the algorithm-independent case manifest.""" return write_json( path, { - "case_schema_version": 1, + "case_schema_version": 2, "cases": [_case_to_dict(case) for case in cases], }, ) diff --git a/scripts/benchmark/motion_generation/config.py b/scripts/benchmark/motion_generation/config.py index 61fb07019..b5f99dcdd 100644 --- a/scripts/benchmark/motion_generation/config.py +++ b/scripts/benchmark/motion_generation/config.py @@ -35,6 +35,7 @@ "FreeSpaceTrackCfg", "PlannerSpecCfg", "ProtocolCfg", + "RobotSpecCfg", "SuiteCfg", "TrackCfg", "load_suite", @@ -84,6 +85,15 @@ class ProtocolCfg: joint_limit_tolerance_rad: float = 1.0e-5 +@configclass +class RobotSpecCfg: + """Robot provider selected for every track in one suite run.""" + + id: str = "franka_panda" + provider: str = "franka_panda" + config: dict[str, Any] = {} + + @configclass class FreeSpaceTrackCfg: """Case matrix for the ``free-space-common`` track.""" @@ -114,6 +124,7 @@ class SuiteCfg: suite_version: str = "free_space_common_v2" profile: str = "smoke" planners: list[PlannerSpecCfg] = [] + robot: RobotSpecCfg = RobotSpecCfg() protocol: ProtocolCfg = ProtocolCfg() tracks: list[TrackCfg] = [] free_space: FreeSpaceTrackCfg = FreeSpaceTrackCfg() @@ -129,6 +140,7 @@ def from_dict(cls, data: dict[str, Any]) -> "SuiteCfg": suite_version=str(data.get("suite_version", "free_space_common_v2")), profile=str(data.get("profile", "smoke")), planners=planners, + robot=RobotSpecCfg(**data.get("robot", {})), protocol=ProtocolCfg(**data.get("protocol", {})), tracks=tracks, free_space=free_space, @@ -171,6 +183,8 @@ def validate_benchmark(self) -> None: "Every planner must define a non-empty id and adapter." ) AlgorithmRole(spec.role) + if not self.robot.id or not self.robot.provider: + raise ValueError("robot must define non-empty id and provider values.") if not self.tracks: raise ValueError("The benchmark suite must declare at least one track.") track_ids = [track.id for track in self.tracks] @@ -199,9 +213,9 @@ def validate_benchmark(self) -> None: _validate_free_space(self.free_space) nmg = next((spec for spec in self.planners if spec.id == "nmg"), None) if nmg is not None: - if float(nmg.config.get("pos_eps", 0.05)) <= 0.0: + if float(nmg.config.get("pos_eps", 0.01)) <= 0.0: raise ValueError("NMG pos_eps must be > 0.") - if float(nmg.config.get("rot_eps", 0.3)) <= 0.0: + if float(nmg.config.get("rot_eps", 0.1)) <= 0.0: raise ValueError("NMG rot_eps must be > 0.") diff --git a/scripts/benchmark/motion_generation/metrics/trajectory.py b/scripts/benchmark/motion_generation/metrics/trajectory.py index fff97a6f1..068323955 100644 --- a/scripts/benchmark/motion_generation/metrics/trajectory.py +++ b/scripts/benchmark/motion_generation/metrics/trajectory.py @@ -42,7 +42,10 @@ def _pose_error_matrices( - waypoints: torch.Tensor, trajectory_poses: torch.Tensor + waypoints: torch.Tensor, + trajectory_poses: torch.Tensor, + *, + rotation_symmetry: str | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Return waypoint-by-sample translation and geodesic rotation errors.""" waypoints = torch.as_tensor(waypoints, dtype=torch.float64) @@ -57,10 +60,31 @@ def _pose_error_matrices( relative = waypoint_rot.transpose(-1, -2) @ trajectory_rot trace = torch.diagonal(relative, dim1=-2, dim2=-1).sum(dim=-1) rotation_error = torch.arccos(torch.clamp((trace - 1.0) * 0.5, -1.0, 1.0)) + if rotation_symmetry == "half_turn_about_z": + symmetric_rot = waypoint_rot.clone() + symmetric_rot[..., :2] = -symmetric_rot[..., :2] + symmetric_relative = symmetric_rot.transpose(-1, -2) @ trajectory_rot + symmetric_trace = torch.diagonal(symmetric_relative, dim1=-2, dim2=-1).sum( + dim=-1 + ) + symmetric_error = torch.arccos( + torch.clamp((symmetric_trace - 1.0) * 0.5, -1.0, 1.0) + ) + rotation_error = torch.minimum(rotation_error, symmetric_error) + elif rotation_symmetry is not None: + raise ValueError( + "rotation_symmetry must be None or 'half_turn_about_z', " + f"got {rotation_symmetry!r}." + ) return pos_error, rotation_error -def get_pose_err(matrix_a: torch.Tensor, matrix_b: torch.Tensor) -> tuple[float, float]: +def get_pose_err( + matrix_a: torch.Tensor, + matrix_b: torch.Tensor, + *, + rotation_symmetry: str | None = None, +) -> tuple[float, float]: """Return translation (m) and geodesic rotation (rad) pose errors.""" tensor_a = torch.as_tensor(matrix_a, dtype=torch.float64) tensor_b = torch.as_tensor(matrix_b, dtype=torch.float64, device=tensor_a.device) @@ -72,6 +96,24 @@ def get_pose_err(matrix_a: torch.Tensor, matrix_b: torch.Tensor) -> tuple[float, relative = tensor_a[:, :3, :3].transpose(-1, -2) @ tensor_b[:, :3, :3] trace = torch.diagonal(relative, dim1=-2, dim2=-1).sum(dim=-1) rotation = torch.arccos(torch.clamp((trace - 1.0) * 0.5, -1.0, 1.0)) + if rotation_symmetry == "half_turn_about_z": + symmetric_b = tensor_b.clone() + symmetric_b[:, :3, :2] = -symmetric_b[:, :3, :2] + symmetric_relative = ( + tensor_a[:, :3, :3].transpose(-1, -2) @ symmetric_b[:, :3, :3] + ) + symmetric_trace = torch.diagonal(symmetric_relative, dim1=-2, dim2=-1).sum( + dim=-1 + ) + symmetric_rotation = torch.arccos( + torch.clamp((symmetric_trace - 1.0) * 0.5, -1.0, 1.0) + ) + rotation = torch.minimum(rotation, symmetric_rotation) + elif rotation_symmetry is not None: + raise ValueError( + "rotation_symmetry must be None or 'half_turn_about_z', " + f"got {rotation_symmetry!r}." + ) return float(translation.mean().item()), float(rotation.mean().item()) @@ -81,6 +123,7 @@ def match_ordered_waypoints( *, position_threshold_m: float, rotation_threshold_rad: float, + rotation_symmetry: str | None = None, ) -> dict[str, object]: """Evaluate ordered arrival and threshold-constrained waypoint errors. @@ -99,9 +142,17 @@ def match_ordered_waypoints( "arrival_indices": [], "position_errors_m": [], "rotation_errors_rad": [], + "min_position_errors_m": [], + "min_rotation_errors_rad": [], + "min_rotation_errors_at_position_rad": [], + "min_position_errors_at_orientation_m": [], } - pos_error, rot_error = _pose_error_matrices(waypoints, trajectory_poses) + pos_error, rot_error = _pose_error_matrices( + waypoints, + trajectory_poses, + rotation_symmetry=rotation_symmetry, + ) arrival_indices: list[int] = [] next_sample = 0 for waypoint_index in range(waypoints.shape[0]): @@ -126,12 +177,35 @@ def match_ordered_waypoints( ] completed = len(arrival_indices) total = int(waypoints.shape[0]) + min_rotation_at_position: list[float | None] = [] + min_position_at_orientation: list[float | None] = [] + for waypoint_index in range(total): + position_hits = pos_error[waypoint_index] <= position_threshold_m + rotation_hits = rot_error[waypoint_index] <= rotation_threshold_rad + min_rotation_at_position.append( + float(rot_error[waypoint_index, position_hits].min().item()) + if bool(position_hits.any().item()) + else None + ) + min_position_at_orientation.append( + float(pos_error[waypoint_index, rotation_hits].min().item()) + if bool(rotation_hits.any().item()) + else None + ) return { "ordered_waypoints_reached": completed == total, "completed_waypoint_ratio": completed / max(total, 1), "arrival_indices": arrival_indices, "position_errors_m": position_errors, "rotation_errors_rad": rotation_errors, + "min_position_errors_m": [ + float(value) for value in pos_error.min(dim=1).values.tolist() + ], + "min_rotation_errors_rad": [ + float(value) for value in rot_error.min(dim=1).values.tolist() + ], + "min_rotation_errors_at_position_rad": min_rotation_at_position, + "min_position_errors_at_orientation_m": min_position_at_orientation, } @@ -141,6 +215,7 @@ def compute_waypoint_errors( *, position_threshold_m: float = 0.01, rotation_threshold_rad: float = 0.1, + rotation_symmetry: str | None = None, ) -> dict[str, float]: """Return ordered, same-sample waypoint errors for one trajectory.""" if isinstance(trajectory_poses, list): @@ -156,6 +231,7 @@ def compute_waypoint_errors( waypoints, position_threshold_m=position_threshold_m, rotation_threshold_rad=rotation_threshold_rad, + rotation_symmetry=rotation_symmetry, ) pos_mm = [float(value) * 1000.0 for value in matched["position_errors_m"]] rot_deg = [ @@ -351,25 +427,45 @@ def compute_case_outcomes( name=control_part, to_matrix=True, ) + native_pose_batch = robot.compute_batch_fk( + qpos=positions, + name=control_part, + to_matrix=True, + ) outcomes: list[CaseOutcome] = [] for env_index in range(case.batch_size): native_qpos = positions[env_index] finite = finite_paths[env_index] if finite: validation_qpos = validation_qpos_batch[env_index] - poses = validation_pose_batch[env_index] + validation_poses = validation_pose_batch[env_index] + # Keep the benchmark's established, sample-count-normalized + # trajectory for all success and path metrics. Native planner + # samples are retained below only for rollout diagnostics. + poses = validation_poses + native_poses = native_pose_batch[env_index] else: validation_qpos = torch.empty( (0, positions.shape[-1]), device=robot.device, dtype=positions.dtype ) + validation_poses = torch.empty( + (0, 4, 4), device=robot.device, dtype=positions.dtype + ) poses = torch.empty((0, 4, 4), device=robot.device, dtype=positions.dtype) + native_poses = torch.empty( + (0, 4, 4), device=robot.device, dtype=positions.dtype + ) waypoints = case.target_waypoints[env_index] + rotation_symmetry = case.case_parameters.get("waypoint_rotation_symmetry") + if rotation_symmetry is not None and not isinstance(rotation_symmetry, str): + raise TypeError("waypoint_rotation_symmetry must be a string or None.") matching = match_ordered_waypoints( poses, waypoints, position_threshold_m=position_threshold_m, rotation_threshold_rad=rotation_threshold_rad, + rotation_symmetry=rotation_symmetry, ) pos_errors_mm = [ float(value) * 1000.0 for value in matching["position_errors_m"] @@ -393,13 +489,30 @@ def compute_case_outcomes( motion_valid = finite and ordered and not joint_violation if poses.shape[0] > 0: - final_pos_m, final_rot_rad = get_pose_err(poses[-1], waypoints[-1]) + final_pos_m, final_rot_rad = get_pose_err( + poses[-1], + waypoints[-1], + rotation_symmetry=rotation_symmetry, + ) + all_pos_error, all_rot_error = _pose_error_matrices( + waypoints, + native_poses, + rotation_symmetry=rotation_symmetry, + ) + min_pos_m = float(all_pos_error[-1].min().item()) + min_rot_rad = float(all_rot_error[-1].min().item()) joint_length, cartesian_length, efficiency = _path_metrics( - validation_qpos, poses, waypoints + validation_qpos, validation_poses, waypoints + ) + trajectory_moved = bool( + torch.linalg.vector_norm(native_qpos - native_qpos[:1], dim=-1).max() + > 1.0e-6 ) else: final_pos_m = final_rot_rad = None + min_pos_m = min_rot_rad = None joint_length = cartesian_length = efficiency = None + trajectory_moved = False if not finite: failure_code = "non_finite_trajectory" @@ -454,6 +567,28 @@ def compute_case_outcomes( path_efficiency=efficiency, failure_code=failure_code, planner_failure_code=planner_failure_code, + min_translation_err_mm=( + min_pos_m * 1000.0 if min_pos_m is not None else None + ), + min_rotation_err_deg=( + min_rot_rad * 180.0 / math.pi if min_rot_rad is not None else None + ), + trajectory_moved=trajectory_moved, + waypoint_min_translation_err_mm=tuple( + float(value) * 1000.0 for value in matching["min_position_errors_m"] + ), + waypoint_min_rotation_err_deg=tuple( + float(value) * 180.0 / math.pi + for value in matching["min_rotation_errors_rad"] + ), + waypoint_min_rotation_err_deg_at_position=tuple( + None if value is None else float(value) * 180.0 / math.pi + for value in matching["min_rotation_errors_at_position_rad"] + ), + waypoint_min_translation_err_mm_at_orientation=tuple( + None if value is None else float(value) * 1000.0 + for value in matching["min_position_errors_at_orientation_m"] + ), ) ) return tuple(outcomes) diff --git a/scripts/benchmark/motion_generation/models.py b/scripts/benchmark/motion_generation/models.py index 02868bc78..24b73756c 100644 --- a/scripts/benchmark/motion_generation/models.py +++ b/scripts/benchmark/motion_generation/models.py @@ -69,7 +69,12 @@ class PlannerMetadata: @dataclass(frozen=True) class BenchmarkCase: - """One env-batched free-space planning input frozen before execution.""" + """One env-batched planner input frozen before execution. + + The first fields retain the free-space case contract. The trailing fields + describe execution/task tracks without forcing planner adapters to know + about a particular robot, atomic skill, or object implementation. + """ suite_version: str track: str @@ -83,6 +88,13 @@ class BenchmarkCase: start_qpos: torch.Tensor target_waypoints: torch.Tensor reference_qpos: torch.Tensor + robot_id: str = "franka_panda" + skill_id: str = "N/A" + object_id: str | None = None + task_difficulty: str | None = None + primary_success: str = "motion_valid" + full_start_qpos: torch.Tensor | None = None + case_parameters: dict[str, object] = field(default_factory=dict) @dataclass(frozen=True) @@ -110,6 +122,21 @@ class CaseOutcome: path_efficiency: float | None failure_code: str | None = None planner_failure_code: str | None = None + execution_success: bool | None = None + task_success: bool | None = None + task_completion_time_s: float | None = None + joint_tracking_rmse_rad: float | None = None + object_lift_delta_m: float | None = None + replan_count: int | None = None + min_translation_err_mm: float | None = None + min_rotation_err_deg: float | None = None + trajectory_moved: bool | None = None + waypoint_min_translation_err_mm: tuple[float, ...] = () + waypoint_min_rotation_err_deg: tuple[float, ...] = () + waypoint_min_rotation_err_deg_at_position: tuple[float | None, ...] = () + waypoint_min_translation_err_mm_at_orientation: tuple[float | None, ...] = () + executed_final_translation_err_mm: float | None = None + executed_final_rotation_err_deg: float | None = None @dataclass(frozen=True) @@ -138,6 +165,15 @@ class TrialRecord: cpu_delta_mb: float | None = None gpu_delta_mb: float | None = None peak_gpu_mb: float | None = None + robot_id: str = "franka_panda" + skill_id: str = "N/A" + object_id: str | None = None + task_difficulty: str | None = None + primary_success: str = "motion_valid" + execution_time_ms: float | None = None + end_to_end_time_ms: float | None = None + trajectory_duration_s: float | None = None + trajectory_waypoints: int | None = None metadata: dict[str, object] = field(default_factory=dict) outcomes: tuple[CaseOutcome, ...] = () diff --git a/scripts/benchmark/motion_generation/planners/__init__.py b/scripts/benchmark/motion_generation/planners/__init__.py index 65b3edabc..27722579f 100644 --- a/scripts/benchmark/motion_generation/planners/__init__.py +++ b/scripts/benchmark/motion_generation/planners/__init__.py @@ -21,13 +21,13 @@ from .base import PlannerAdapter, PlannerContext from .curobo import CuroboAdapter from .ik_interpolate import IkInterpolateAdapter -from .neural import NeuralAdapterStub +from .nmg_onnx import NmgOnnxAdapter from .toppra import ToppraAdapter __all__ = [ "CuroboAdapter", "IkInterpolateAdapter", - "NeuralAdapterStub", + "NmgOnnxAdapter", "PlannerAdapter", "PlannerContext", "ToppraAdapter", diff --git a/scripts/benchmark/motion_generation/planners/base.py b/scripts/benchmark/motion_generation/planners/base.py index 895e5f11b..534796ec0 100644 --- a/scripts/benchmark/motion_generation/planners/base.py +++ b/scripts/benchmark/motion_generation/planners/base.py @@ -31,6 +31,7 @@ if TYPE_CHECKING: from embodichain.lab.sim.objects import Robot + from embodichain.lab.sim.planners import MotionGenerator __all__ = ["PlannerAdapter", "PlannerContext"] @@ -43,6 +44,7 @@ class PlannerContext: control_part: str device: torch.device sample_interval: int + robot_id: str = "unknown" class PlannerAdapter(ABC): @@ -66,12 +68,15 @@ def metadata(self) -> PlannerMetadata: adapter=self.spec.adapter, config_hash=stable_hash(self.spec.config), capabilities=self.capabilities, - model_revision=str( - self.spec.config.get("model_revision", self.model_revision) - ), + model_revision=self._resolved_model_revision(), + supported_robots=(self.context.robot_id,), parameters=dict(self.spec.config), ) + def _resolved_model_revision(self) -> str: + """Resolve the recorded model identity for this planner instance.""" + return str(self.spec.config.get("model_revision", self.model_revision)) + def availability(self) -> tuple[bool, str | None]: """Return whether this adapter can run in the current process.""" return True, None @@ -84,6 +89,25 @@ def prepare(self, case: BenchmarkCase) -> dict[str, object] | None: """Prepare a lazy backend, or return ``None`` when not applicable.""" return None + @property + def motion_policy_planner(self) -> str: + """Return the backend name pinned into Atomic Action motion policies.""" + return self.spec.adapter + + def require_motion_generator(self) -> "MotionGenerator": + """Return the adapter-owned MotionGenerator or fail clearly. + + Atomic-task scenarios use this boundary instead of reaching into a + backend implementation. Every planner that opts into the + ``atomic_action`` capability must expose its generator here. + """ + motion_generator = getattr(self, "motion_generator", None) + if motion_generator is None: + raise RuntimeError( + f"Planner adapter {self.spec.id!r} does not expose a MotionGenerator." + ) + return motion_generator + @abstractmethod def plan(self, case: BenchmarkCase) -> PlanResult: """Plan one env-batched benchmark case.""" diff --git a/scripts/benchmark/motion_generation/planners/curobo.py b/scripts/benchmark/motion_generation/planners/curobo.py index 2b4e34716..30ff1c1dd 100644 --- a/scripts/benchmark/motion_generation/planners/curobo.py +++ b/scripts/benchmark/motion_generation/planners/curobo.py @@ -46,7 +46,15 @@ class CuroboAdapter(PlannerAdapter): """Run cuRobo with a frozen, empty-world operational configuration.""" - capabilities = frozenset({"eef_waypoint", "batched", "empty_world"}) + capabilities = frozenset( + { + "eef_waypoint", + "joint_waypoint", + "batched", + "empty_world", + "atomic_action", + } + ) model_revision = "curobo-v2" separate_prepare = True @@ -69,7 +77,7 @@ def build(self) -> None: auto_values = dict(values.get("auto_gen", {})) if bool(world_values.get("multi_env", False)): raise ValueError( - "free-space-common requires one shared empty cuRobo world " + "The current cuRobo benchmark adapter requires one shared empty world " "with world.multi_env=false." ) world = CuroboWorldCfg( diff --git a/scripts/benchmark/motion_generation/planners/neural.py b/scripts/benchmark/motion_generation/planners/neural.py deleted file mode 100644 index 20af60396..000000000 --- a/scripts/benchmark/motion_generation/planners/neural.py +++ /dev/null @@ -1,55 +0,0 @@ -# ---------------------------------------------------------------------------- -# 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. -# ---------------------------------------------------------------------------- - -"""Capability-aware NMG adapter stub reserved for a future checkpoint.""" - -from __future__ import annotations - -from embodichain.lab.sim.planners import PlanResult - -from ..models import BenchmarkCase -from ..registry import register_planner_adapter -from .base import PlannerAdapter - -__all__ = ["NeuralAdapterStub"] - - -class NeuralAdapterStub(PlannerAdapter): - """Expose configurable NMG precision without initializing an unavailable model.""" - - capabilities = frozenset({"eef_waypoint", "batched", "empty_world"}) - model_revision = "not-ready" - - def availability(self) -> tuple[bool, str | None]: - """Mark the placeholder unsupported until the checkpoint contract lands.""" - pos_eps = float(self.spec.config.get("pos_eps", 0.05)) - rot_eps = float(self.spec.config.get("rot_eps", 0.3)) - return ( - False, - "NMG adapter is a stub pending the production checkpoint; " - f"configured pos_eps={pos_eps} m, rot_eps={rot_eps} rad.", - ) - - def build(self) -> None: - """Reject accidental construction of the explicit placeholder.""" - raise RuntimeError("The NMG adapter is not implemented yet.") - - def plan(self, case: BenchmarkCase) -> PlanResult: # noqa: ARG002 - """Reject accidental execution of the explicit placeholder.""" - raise RuntimeError("The NMG adapter is not implemented yet.") - - -register_planner_adapter("neural_stub", NeuralAdapterStub) diff --git a/scripts/benchmark/motion_generation/planners/nmg_onnx.py b/scripts/benchmark/motion_generation/planners/nmg_onnx.py new file mode 100644 index 000000000..6210a8f2a --- /dev/null +++ b/scripts/benchmark/motion_generation/planners/nmg_onnx.py @@ -0,0 +1,136 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""NMG ONNX benchmark adapter.""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path + +from embodichain.lab.sim.planners import ( + MotionGenCfg, + MotionGenOptions, + MotionGenerator, + NeuralPlanOptions, + NeuralPlannerCfg, + PlanResult, + PlanState, +) + +from ..config import PlannerSpecCfg +from ..models import BenchmarkCase +from ..registry import register_planner_adapter +from .base import PlannerAdapter, PlannerContext + +__all__ = ["NmgOnnxAdapter"] + + +class NmgOnnxAdapter(PlannerAdapter): + """Run a standalone NMG ONNX policy through :class:`NeuralPlanner`.""" + + capabilities = frozenset( + { + "eef_waypoint", + "joint_waypoint", + "batched", + "empty_world", + "atomic_action", + } + ) + + def __init__(self, spec: PlannerSpecCfg, context: PlannerContext) -> None: + super().__init__(spec, context) + self.motion_generator: MotionGenerator | None = None + + def _resolved_model_revision(self) -> str: + """Use an explicit revision or derive one from the runtime ONNX path.""" + configured = self.spec.config.get("model_revision") + if configured: + return str(configured) + model_path = self.spec.config.get("onnx_model_path") + if model_path: + return Path(str(model_path)).expanduser().stem + return super()._resolved_model_revision() + + def availability(self) -> tuple[bool, str | None]: + """Require ONNX Runtime and a resolved standalone policy path.""" + if importlib.util.find_spec("onnxruntime") is None: + return False, "onnxruntime is not installed." + model_path = self.spec.config.get("onnx_model_path") + if not model_path: + return False, "NMG requires config.onnx_model_path or --nmg-onnx-path." + path = Path(str(model_path)).expanduser() + if not path.is_file(): + return False, f"NMG ONNX policy does not exist: {path}." + if path.suffix.lower() != ".onnx": + return False, f"NMG requires a standalone .onnx policy, got {path}." + return True, None + + @property + def motion_policy_planner(self) -> str: + """Select EmbodiChain's neural MotionGenerator backend.""" + return "neural" + + def build(self) -> None: + """Construct the ONNX-backed MotionGenerator.""" + values = self.spec.config + model_path = str(Path(str(values["onnx_model_path"])).expanduser()) + providers = values.get("onnx_providers") + planner_cfg = NeuralPlannerCfg( + robot_uid=self.context.robot.uid, + onnx_model_path=model_path, + control_part=self.context.control_part, + max_steps=int(values.get("max_steps", 240)), + action_scale=float(values.get("action_scale", 0.2)), + num_arm_joints=int(values.get("num_arm_joints", 7)), + num_waypoints=int(values.get("num_waypoints", 8)), + use_relative_obs=bool(values.get("use_relative_obs", True)), + canonicalize_quat_obs=bool(values.get("canonicalize_quat_obs", True)), + intermediate_orientation=bool(values.get("intermediate_orientation", True)), + pos_eps=float(values.get("pos_eps", 0.01)), + rot_eps=float(values.get("rot_eps", 0.1)), + joint_eps=float(values.get("joint_eps", 0.02)), + onnx_providers=list(providers) if providers is not None else None, + policy_frame_from_world=values.get("policy_frame_from_world"), + runtime_tcp_from_policy_tcp=values.get("runtime_tcp_from_policy_tcp"), + dt=float(values.get("dt", 0.01)), + ) + self.motion_generator = MotionGenerator(MotionGenCfg(planner_cfg=planner_cfg)) + + def plan(self, case: BenchmarkCase) -> PlanResult: + """Plan all ordered Cartesian constraints in one closed-loop rollout.""" + if self.motion_generator is None: + raise RuntimeError("NMG adapter must be built before plan().") + targets = [ + PlanState.from_xpos(case.target_waypoints[:, index]) + for index in range(case.num_waypoints) + ] + return self.motion_generator.generate( + targets, + MotionGenOptions( + start_qpos=case.start_qpos, + control_part=self.context.control_part, + plan_opts=NeuralPlanOptions(), + ), + ) + + def close(self) -> None: + """Release the adapter-owned MotionGenerator.""" + self._close_motion_generator() + + +register_planner_adapter("nmg_onnx", NmgOnnxAdapter) diff --git a/scripts/benchmark/motion_generation/registry.py b/scripts/benchmark/motion_generation/registry.py index d10efb3fe..7611b015e 100644 --- a/scripts/benchmark/motion_generation/registry.py +++ b/scripts/benchmark/motion_generation/registry.py @@ -21,22 +21,27 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from .config import PlannerSpecCfg + from .config import PlannerSpecCfg, RobotSpecCfg from .planners.base import PlannerAdapter, PlannerContext + from .robots.base import RobotProvider from .scenarios.base import ScenarioProvider __all__ = [ "create_planner_adapter", "create_scenario_provider", + "create_robot_provider", "planner_adapter_names", "register_planner_adapter", "register_scenario_provider", + "register_robot_provider", + "robot_provider_names", "scenario_provider_names", "unregister_planner_adapter", ] _PLANNER_ADAPTERS: dict[str, type["PlannerAdapter"]] = {} _SCENARIO_PROVIDERS: dict[str, type["ScenarioProvider"]] = {} +_ROBOT_PROVIDERS: dict[str, type["RobotProvider"]] = {} def register_planner_adapter(name: str, adapter_cls: type["PlannerAdapter"]) -> None: @@ -73,6 +78,33 @@ def create_planner_adapter( return adapter_cls(spec=spec, context=context) +def register_robot_provider(name: str, provider_cls: type["RobotProvider"]) -> None: + """Register one robot provider under a stable suite name.""" + if not name: + raise ValueError("Robot provider name must not be empty.") + previous = _ROBOT_PROVIDERS.get(name) + if previous is not None and previous is not provider_cls: + raise ValueError(f"Robot provider {name!r} is already registered.") + _ROBOT_PROVIDERS[name] = provider_cls + + +def robot_provider_names() -> tuple[str, ...]: + """Return registered robot-provider names in deterministic order.""" + return tuple(sorted(_ROBOT_PROVIDERS)) + + +def create_robot_provider(spec: "RobotSpecCfg") -> "RobotProvider": + """Construct the robot provider selected by a suite.""" + try: + provider_cls = _ROBOT_PROVIDERS[spec.provider] + except KeyError as exc: + raise ValueError( + f"Unknown robot provider {spec.provider!r}; " + f"registered providers: {robot_provider_names()}." + ) from exc + return provider_cls(spec) + + def register_scenario_provider( name: str, provider_cls: type["ScenarioProvider"] ) -> None: diff --git a/scripts/benchmark/motion_generation/reporting.py b/scripts/benchmark/motion_generation/reporting.py index e5d26af63..49fc50f75 100644 --- a/scripts/benchmark/motion_generation/reporting.py +++ b/scripts/benchmark/motion_generation/reporting.py @@ -14,7 +14,7 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Render the free-space benchmark as exactly three Markdown tables.""" +"""Render planner and Atomic Task tracks as exactly three Markdown tables.""" from __future__ import annotations @@ -30,6 +30,10 @@ "track", "algorithm", "algorithm_role", + "robot", + "skill", + "object", + "task_difficulty", "batch_size", "waypoint_count", "num_trials", @@ -45,6 +49,10 @@ "cpu_delta_mb", "gpu_delta_mb", "peak_gpu_mb", + "execution_time_ms", + "end_to_end_time_ms", + "trajectory_duration_s", + "trajectory_waypoints", ) METRIC_COLUMNS = ( @@ -52,6 +60,11 @@ "scenario", "algorithm", "algorithm_role", + "robot", + "skill", + "object", + "task_difficulty", + "primary_success", "batch_size", "waypoint_count", "path_shape", @@ -61,6 +74,9 @@ "coverage_rate", "success_rate", "planning_success_rate", + "motion_valid_rate", + "execution_success_rate", + "task_success_rate", "ordered_waypoint_success_rate", "waypoint_completion_rate", "final_pos_err_mm", @@ -71,6 +87,10 @@ "joint_path_length_rad", "cartesian_path_length_m", "path_efficiency", + "task_completion_time_s", + "joint_tracking_rmse_rad", + "object_lift_delta_m", + "replan_count", "top_failure", ) @@ -86,6 +106,7 @@ "overall_success_rate", "planning_success_rate", "motion_valid_rate", + "execution_success_rate", "task_success_rate", "latency_p95_ms", "peak_gpu_mb", @@ -132,7 +153,7 @@ def write_markdown_report( output = Path(path) output.parent.mkdir(parents=True, exist_ok=True) lines = [ - "# Motion Generation Benchmark Report", + "# Planner Motion Generation & Atomic Skill Benchmark Report", "", f"Generated at: {datetime.now(timezone.utc).isoformat(timespec='seconds')}", "", @@ -158,7 +179,8 @@ def write_markdown_report( "## Success & Other Metrics", "", "Continuous error/path columns are conditioned on `motion_valid` " - "outcomes; use `n_valid` as the denominator before comparing them.", + "outcomes; use `n_valid` as the denominator. Atomic Task " + "`success_rate` follows the case-owned `primary_success` stage.", "", ] ) diff --git a/scripts/benchmark/motion_generation/robots/__init__.py b/scripts/benchmark/motion_generation/robots/__init__.py new file mode 100644 index 000000000..417f45245 --- /dev/null +++ b/scripts/benchmark/motion_generation/robots/__init__.py @@ -0,0 +1,24 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Built-in robot providers for motion-generation benchmarks.""" + +from __future__ import annotations + +from .base import RobotProvider +from .franka import FrankaPandaProvider, FrankaPgiProvider + +__all__ = ["FrankaPandaProvider", "FrankaPgiProvider", "RobotProvider"] diff --git a/scripts/benchmark/motion_generation/robots/base.py b/scripts/benchmark/motion_generation/robots/base.py new file mode 100644 index 000000000..5d7e1786f --- /dev/null +++ b/scripts/benchmark/motion_generation/robots/base.py @@ -0,0 +1,48 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Robot-provider contract for planner benchmark suites.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from embodichain.lab.sim import SimulationManager + from embodichain.lab.sim.cfg import RobotCfg + from embodichain.lab.sim.objects import Robot + + from ..config import RobotSpecCfg + +__all__ = ["RobotProvider"] + + +class RobotProvider(ABC): + """Build one benchmark embodiment behind a stable suite identifier.""" + + control_part: str = "arm" + + def __init__(self, spec: "RobotSpecCfg") -> None: + self.spec = spec + + @abstractmethod + def build_cfg(self) -> "RobotCfg": + """Build the robot configuration without mutating a simulation.""" + + def add_robot(self, simulation: "SimulationManager") -> "Robot": + """Add the configured robot to a simulation.""" + return simulation.add_robot(cfg=self.build_cfg()) diff --git a/scripts/benchmark/motion_generation/robots/franka.py b/scripts/benchmark/motion_generation/robots/franka.py new file mode 100644 index 000000000..e7cf6280c --- /dev/null +++ b/scripts/benchmark/motion_generation/robots/franka.py @@ -0,0 +1,68 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Built-in Franka robot providers.""" + +from __future__ import annotations + +from embodichain.lab.sim.cfg import RobotCfg +from embodichain.lab.sim.robots import FrankaPandaCfg +from embodichain.lab.sim.utility.cfg_utils import merge_robot_cfg +from scripts.tutorials.atomic_action.tutorial_utils import ( + create_franka_panda_robot_cfg, +) + +from ..registry import register_robot_provider +from .base import RobotProvider + +__all__ = ["FrankaPandaProvider", "FrankaPgiProvider"] + + +class FrankaPandaProvider(RobotProvider): + """Build the stock Franka Panda used by the free-space track.""" + + def build_cfg(self) -> RobotCfg: + """Build a stock Panda while applying optional suite overrides.""" + values = { + "uid": "benchmark_franka_panda", + "robot_type": "panda", + **dict(self.spec.config), + } + return FrankaPandaCfg.from_dict(values) + + +class FrankaPgiProvider(RobotProvider): + """Build a Franka arm assembled with the tutorial PGI gripper. + + The configuration mirrors the Franka compatibility in + ``scripts/tutorials/atomic_action``: an arm-only Panda URDF, the shared + ``DH_PGI_140_80`` component, a 180-degree base rotation, and the PGI TCP. + """ + + def build_cfg(self) -> RobotCfg: + """Build the Franka + PGI benchmark configuration.""" + cfg = create_franka_panda_robot_cfg() + return merge_robot_cfg( + cfg, + { + "uid": "benchmark_franka_pgi", + **dict(self.spec.config), + }, + ) + + +register_robot_provider("franka_panda", FrankaPandaProvider) +register_robot_provider("franka_pgi", FrankaPgiProvider) diff --git a/scripts/benchmark/motion_generation/run_benchmark.py b/scripts/benchmark/motion_generation/run_benchmark.py index 4a9ff9ee6..541d50eeb 100644 --- a/scripts/benchmark/motion_generation/run_benchmark.py +++ b/scripts/benchmark/motion_generation/run_benchmark.py @@ -14,13 +14,14 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Run the extensible free-space motion-generation benchmark. +"""Run the extensible planner motion-generation benchmark. cuRobo is the default primary baseline. IK interpolation and TOPPRA are -optional diagnostic baselines. NMG remains an explicitly configurable, -unsupported adapter stub until its production checkpoint contract is ready. +optional diagnostic baselines. NMG is loaded from a standalone ONNX policy. -Run: ``embodichain benchmark motion-generation --suite smoke`` +Run: ``python -m scripts.benchmark.motion_generation.run_benchmark --suite +smoke`` or select the Franka + PGI Atomic Task slice with +``--suite atomic_franka_pgi_curobo``. """ from __future__ import annotations @@ -31,6 +32,7 @@ from typing import TYPE_CHECKING from .config import PlannerSpecCfg, SuiteCfg, load_suite +from .video import VideoRecordCfg, video_cfg_from_args if TYPE_CHECKING: from .runner import BenchmarkRunResult @@ -43,11 +45,15 @@ def add_parser_arguments(parser: argparse.ArgumentParser) -> None: - """Add free-space benchmark options to an existing argument parser.""" + """Add planner benchmark options to an existing argument parser.""" parser.add_argument( "--suite", default="smoke", - help="Suite short name (smoke/coverage) or an explicit YAML path.", + help=( + "Suite short name (smoke/coverage/atomic_franka_pgi_curobo/" + "atomic_franka_pgi_curobo_randomized) " + "or an explicit YAML path." + ), ) parser.add_argument( "--algorithms", @@ -92,9 +98,9 @@ def add_parser_arguments(parser: argparse.ArgumentParser) -> None: help="NMG internal waypoint rotation threshold in radians.", ) parser.add_argument( - "--checkpoint-path", + "--nmg-onnx-path", default=None, - help="Reserved NMG checkpoint path; the current NMG adapter remains a stub.", + help="Path to the standalone NMG ONNX policy.", ) parser.add_argument( "--output-root", default="outputs/benchmarks", help="Artifact root directory." @@ -105,6 +111,36 @@ def add_parser_arguments(parser: argparse.ArgumentParser) -> None: parser.add_argument( "--no-headless", action="store_false", dest="headless", help="Open a viewer." ) + parser.add_argument( + "--record-video", + action="store_true", + help="Record Atomic Task measured physics-replay videos after evaluation.", + ) + parser.add_argument( + "--record-failed-video", + action="store_true", + help="With --record-video, also record failed cases as static debug scenes.", + ) + parser.add_argument( + "--video-case-limit", + type=int, + default=0, + help="Maximum recorded videos. Use 0 to record every selected case.", + ) + parser.add_argument( + "--video-dir", + default=None, + help="Override video directory. Default is /videos.", + ) + parser.add_argument("--video-fps", type=int, default=20) + parser.add_argument("--video-width", type=int, default=640) + parser.add_argument("--video-height", type=int, default=480) + parser.add_argument( + "--video-max-memory", + type=int, + default=2048, + help="Maximum recorder frame-buffer memory in MB.", + ) def _planner_by_id(suite: SuiteCfg, planner_id: str) -> PlannerSpecCfg: @@ -154,7 +190,7 @@ def _apply_overrides( rotation_threshold_rad: float | None = None, nmg_pos_eps: float | None = None, nmg_rot_eps: float | None = None, - checkpoint_path: str | None = None, + nmg_onnx_path: str | None = None, ) -> None: """Apply explicit CLI/programmatic overrides to a loaded suite.""" if batch_sizes is not None: @@ -167,6 +203,9 @@ def _apply_overrides( suite.free_space.start_state_bins = start_state_bins if seeds is not None: suite.free_space.seeds = seeds + for track in suite.tracks: + if track.scenario != "free_space" and "seeds" in track.config: + track.config["seeds"] = list(seeds) if num_trials is not None: suite.protocol.measured_trials = num_trials if warmup_trials is not None: @@ -186,8 +225,8 @@ def _apply_overrides( nmg.config["pos_eps"] = nmg_pos_eps if nmg_rot_eps is not None: nmg.config["rot_eps"] = nmg_rot_eps - if checkpoint_path is not None: - nmg.config["checkpoint_path"] = str(Path(checkpoint_path)) + if nmg_onnx_path is not None: + nmg.config["onnx_model_path"] = str(Path(nmg_onnx_path)) suite.validate_benchmark() @@ -195,7 +234,7 @@ def run_all_benchmarks( num_waypoints_list: list[int] | None = None, sim_device: str = "auto", headless: bool = True, - checkpoint_path: str | None = None, + nmg_onnx_path: str | None = None, *, suite_name: str = "smoke", algorithms: list[str] | None = None, @@ -213,8 +252,9 @@ def run_all_benchmarks( nmg_pos_eps: float | None = None, nmg_rot_eps: float | None = None, output_root: str | Path = "outputs/benchmarks", + video: VideoRecordCfg | None = None, ) -> BenchmarkRunResult: - """Resolve configuration and run all selected free-space benchmarks.""" + """Resolve configuration and run all selected benchmark tracks.""" from .runner import BenchmarkRunner suite = load_suite(suite_name) @@ -233,7 +273,7 @@ def run_all_benchmarks( rotation_threshold_rad=rotation_threshold_rad, nmg_pos_eps=nmg_pos_eps, nmg_rot_eps=nmg_rot_eps, - checkpoint_path=checkpoint_path, + nmg_onnx_path=nmg_onnx_path, ) specs = _resolve_planners(suite, algorithms, list(extra_baselines or [])) return BenchmarkRunner( @@ -242,6 +282,7 @@ def run_all_benchmarks( device=sim_device, headless=headless, output_root=output_root, + video=video, ).run() @@ -251,7 +292,7 @@ def run_from_args(args: argparse.Namespace) -> BenchmarkRunResult: num_waypoints_list=args.num_waypoints, sim_device=args.device, headless=args.headless, - checkpoint_path=args.checkpoint_path, + nmg_onnx_path=args.nmg_onnx_path, suite_name=args.suite, algorithms=args.algorithms, extra_baselines=args.extra_baselines, @@ -268,13 +309,14 @@ def run_from_args(args: argparse.Namespace) -> BenchmarkRunResult: nmg_pos_eps=args.nmg_pos_eps, nmg_rot_eps=args.nmg_rot_eps, output_root=args.output_root, + video=video_cfg_from_args(args), ) def _parse_args() -> argparse.Namespace: """Parse standalone module arguments using the unified option schema.""" parser = argparse.ArgumentParser( - description="Benchmark motion generation on fixed free-space cases." + description="Benchmark planners on fixed motion and Atomic Task cases." ) add_parser_arguments(parser) return parser.parse_args() diff --git a/scripts/benchmark/motion_generation/runner.py b/scripts/benchmark/motion_generation/runner.py index 6f6badc87..c1ce0dd70 100644 --- a/scripts/benchmark/motion_generation/runner.py +++ b/scripts/benchmark/motion_generation/runner.py @@ -25,10 +25,9 @@ import torch from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.planners.utils import PlanResult -from embodichain.lab.sim.robots import FrankaPandaCfg from . import planners as _builtin_planners # noqa: F401 - registry side effects +from . import robots as _builtin_robots # noqa: F401 - registry side effects from . import scenarios as _builtin_scenarios # noqa: F401 - registry side effects from .aggregation import aggregate_results from .artifacts import ( @@ -40,8 +39,7 @@ write_resolved_suite, ) from .config import PlannerSpecCfg, SuiteCfg -from .metrics import compute_case_outcomes, timed_call -from .metrics.trajectory import make_failure_outcomes +from .metrics import timed_call from .models import ( BenchmarkCase, PlannerMetadata, @@ -49,8 +47,19 @@ TrialRecord, ) from .planners.base import PlannerAdapter, PlannerContext -from .registry import create_planner_adapter, create_scenario_provider +from .registry import ( + create_planner_adapter, + create_robot_provider, + create_scenario_provider, +) from .reporting import write_markdown_report +from .scenarios.base import ScenarioEvaluation, ScenarioProvider +from .scenarios.free_space import FreeSpaceScenario +from .video import ( + VideoRecordCfg, + should_record_case, + summarize_video_recording, +) if TYPE_CHECKING: from collections.abc import Callable @@ -60,8 +69,6 @@ __all__ = ["BenchmarkRunResult", "BenchmarkRunner", "resolve_device"] _T = TypeVar("_T") -_CONTROL_PART = "arm" -_ROBOT_UID = "benchmark_franka_panda" @dataclass(frozen=True) @@ -103,19 +110,25 @@ def __init__( device: str = "auto", headless: bool = True, output_root: str | Path = "outputs/benchmarks", + video: VideoRecordCfg | None = None, ) -> None: self.suite = suite self.planner_specs = planner_specs self.device = resolve_device(device) self.headless = headless self.output_root = Path(output_root) + self.video = VideoRecordCfg() if video is None else video + self.robot_provider = create_robot_provider(suite.robot) + self.control_part = self.robot_provider.control_part self.records: list[TrialRecord] = [] self.cases: list[BenchmarkCase] = [] self.metadata: dict[str, PlannerMetadata] = {} self.notes: list[str] = [] + self._run_dir: Path | None = None + self._video_paths: list[str] = [] def _create_simulation(self, batch_size: int) -> tuple[SimulationManager, "Robot"]: - """Create one isolated Franka simulator for a fixed batch size.""" + """Create one isolated suite-selected robot for a fixed batch size.""" sim = SimulationManager( SimulationManagerCfg( headless=self.headless, @@ -124,22 +137,10 @@ def _create_simulation(self, batch_size: int) -> tuple[SimulationManager, "Robot arena_space=2.0, ) ) - robot = sim.add_robot( - cfg=FrankaPandaCfg.from_dict({"uid": _ROBOT_UID, "robot_type": "panda"}) - ) + robot = self.robot_provider.add_robot(sim) sim.update(step=1) return sim, robot - @staticmethod - def _set_case_start( - sim: SimulationManager, robot: "Robot", case: BenchmarkCase - ) -> None: - """Restore current and target robot state outside the timed region.""" - robot.set_qpos(case.start_qpos, name=_CONTROL_PART, target=False) - robot.set_qpos(case.start_qpos, name=_CONTROL_PART, target=True) - robot.clear_dynamics() - sim.update(step=1) - def _append(self, writer: TrialJsonlWriter, record: TrialRecord) -> None: """Retain and immediately persist one raw record.""" self.records.append(record) @@ -169,6 +170,11 @@ def _base_record( "waypoint_count": case.num_waypoints, "path_shape": case.path_shape, "start_state_bin": case.start_state_bin, + "robot_id": case.robot_id, + "skill_id": case.skill_id, + "object_id": case.object_id, + "task_difficulty": case.task_difficulty, + "primary_success": case.primary_success, "phase": phase, } @@ -237,49 +243,57 @@ def _run_plan_call( adapter: PlannerAdapter, metadata: PlannerMetadata, case: BenchmarkCase, + provider: ScenarioProvider, phase: TrialPhase, repeat: int, ) -> None: """Time one plan, validate outside timing, and persist the record.""" - self._set_case_start(sim, robot, case) - measured = timed_call(lambda: _capture(lambda: adapter.plan(case))) + provider.reset_case(sim, robot, case, self.control_part) + measured = timed_call( + lambda: _capture(lambda: provider.plan_case(adapter, case)) + ) result, error = measured.result failure_code = None failure_message = None status = "ok" + evaluation: ScenarioEvaluation | None = None if error is not None: status = "error" failure_code = "planner_exception" failure_message = str(error) - outcomes = make_failure_outcomes(case.batch_size, failure_code) - elif not isinstance(result, PlanResult): + outcomes = provider.failure_outcomes(case, failure_code) + elif (contract_error := provider.plan_contract_error(result)) is not None: status = "error" failure_code = "planner_contract_error" - failure_message = f"Expected PlanResult, got {type(result).__name__}." - outcomes = make_failure_outcomes(case.batch_size, failure_code) + failure_message = contract_error + outcomes = provider.failure_outcomes(case, failure_code) elif phase in (TrialPhase.WARMUP, TrialPhase.COLD): # Cold/warmup timing must not pay for FK validation that is unused # by aggregation. outcomes = () else: try: - outcomes = compute_case_outcomes( + evaluation = provider.evaluate_case( result, case, robot, - _CONTROL_PART, - validation_samples=self.suite.protocol.validation_samples, - position_threshold_m=self.suite.protocol.position_threshold_m, - rotation_threshold_rad=self.suite.protocol.rotation_threshold_rad, - joint_limit_tolerance_rad=( - self.suite.protocol.joint_limit_tolerance_rad - ), + self.control_part, + self.suite, + planning_time_ms=measured.cost_time_ms, ) + outcomes = evaluation.outcomes except Exception as exc: # noqa: BLE001 - metric failure is recorded status = "error" failure_code = "metric_evaluation_error" failure_message = str(exc) - outcomes = make_failure_outcomes(case.batch_size, failure_code) + outcomes = provider.failure_outcomes(case, failure_code) + + record_metadata = {} if evaluation is None else dict(evaluation.metadata) + video_path = self._maybe_record_replay( + provider, result, case, evaluation, metadata.algorithm_id, phase + ) + if video_path is not None: + record_metadata["video_path"] = str(video_path) self._append( writer, @@ -292,17 +306,68 @@ def _run_plan_call( cpu_delta_mb=measured.cpu_delta_mb, gpu_delta_mb=measured.gpu_delta_mb, peak_gpu_mb=measured.peak_gpu_mb, + execution_time_ms=( + None if evaluation is None else evaluation.execution_time_ms + ), + end_to_end_time_ms=( + None if evaluation is None else evaluation.end_to_end_time_ms + ), + trajectory_duration_s=( + None if evaluation is None else evaluation.trajectory_duration_s + ), + trajectory_waypoints=( + None if evaluation is None else evaluation.trajectory_waypoints + ), + metadata=record_metadata, outcomes=outcomes, ), ) if phase is not TrialPhase.WARMUP: print( f" {metadata.algorithm_id:<16} B={case.batch_size:>3d} " - f"W={case.num_waypoints} {case.path_shape:<16} " + f"W={case.num_waypoints} {case.skill_id:<20} " f"{phase.value:<8} {measured.cost_time_ms:>10.3f} ms " f"status={status}" ) + def _maybe_record_replay( + self, + provider: ScenarioProvider, + result: object, + case: BenchmarkCase, + evaluation: ScenarioEvaluation | None, + algorithm_id: str, + phase: TrialPhase, + ) -> Path | None: + """Record one measured Atomic Task replay outside planner timing.""" + if phase is not TrialPhase.MEASURED or not self.video.enabled: + return None + success = bool( + evaluation is not None + and evaluation.outcomes + and all(bool(outcome.task_success) for outcome in evaluation.outcomes) + ) + if not should_record_case(self.video, len(self._video_paths), success): + return None + if self._run_dir is None: + raise RuntimeError("Benchmark run directory is not initialized.") + output_dir = ( + self.video.output_dir + if self.video.output_dir is not None + else self._run_dir / "videos" + ) + video_path = provider.record_replay( + result, + case, + evaluation, + output_dir=output_dir, + algorithm_id=algorithm_id, + video=self.video, + ) + if video_path is not None: + self._video_paths.append(str(video_path)) + return video_path + def _run_adapter( self, writer: TrialJsonlWriter, @@ -311,13 +376,16 @@ def _run_adapter( spec: PlannerSpecCfg, cases: list[BenchmarkCase], required_capabilities: frozenset[str], + provider: ScenarioProvider | None = None, ) -> None: """Execute one adapter over every case for a fixed simulator batch.""" + provider = provider or FreeSpaceScenario() context = PlannerContext( robot=robot, - control_part=_CONTROL_PART, + control_part=self.control_part, device=self.device, sample_interval=self.suite.protocol.sample_interval, + robot_id=self.suite.robot.id, ) adapter = create_planner_adapter(spec, context) metadata = adapter.metadata @@ -354,6 +422,7 @@ def _run_adapter( if build_error is not None: adapter.close() return + scenario_prepared = False try: if adapter.separate_prepare: _, prepare_error = self._record_timed_lifecycle( @@ -366,6 +435,25 @@ def _run_adapter( if prepare_error is not None: return + try: + provider.prepare_planner(adapter, first_case) + scenario_prepared = True + except Exception as exc: # noqa: BLE001 - recorded benchmark failure + self._append( + writer, + TrialRecord( + **self._base_record(metadata, first_case, TrialPhase.PREPARE), + status="error", + failure_code="scenario_prepare_error", + failure_message=str(exc), + ), + ) + self.notes.append( + f"{metadata.algorithm_id} scenario prepare failed for " + f"B={first_case.batch_size}: {exc}" + ) + return + self._run_plan_call( writer, sim, @@ -373,6 +461,7 @@ def _run_adapter( adapter, metadata, first_case, + provider, TrialPhase.COLD, repeat=-1, ) @@ -385,6 +474,7 @@ def _run_adapter( adapter, metadata, case, + provider, TrialPhase.WARMUP, repeat=warmup_index, ) @@ -396,15 +486,19 @@ def _run_adapter( adapter, metadata, case, + provider, TrialPhase.MEASURED, repeat=repeat, ) finally: + if scenario_prepared: + provider.close_planner(adapter) adapter.close() def run(self) -> BenchmarkRunResult: """Run the suite and write all required artifacts.""" run_dir = create_run_directory(self.output_root, self.suite.name) + self._run_dir = run_dir write_resolved_suite(run_dir / "resolved_suite.yaml", self.suite) write_json(run_dir / "environment.json", environment_metadata()) writer = TrialJsonlWriter(run_dir / "trials.jsonl") @@ -423,10 +517,19 @@ def run(self) -> BenchmarkRunResult: provider = create_scenario_provider(track.scenario) for batch_size in provider.batch_sizes(self.suite, track): sim: SimulationManager | None = None + runtime_configured = False try: sim, robot = self._create_simulation(batch_size) + provider.configure_runtime( + sim, + robot, + self.suite, + track, + self.control_part, + ) + runtime_configured = True cases = provider.generate_cases( - self.suite, track, robot, _CONTROL_PART, batch_size + self.suite, track, robot, self.control_part, batch_size ) self.cases.extend(cases) for spec in self.planner_specs: @@ -437,8 +540,11 @@ def run(self) -> BenchmarkRunResult: spec, cases, provider.required_capabilities, + provider, ) finally: + if runtime_configured: + provider.close_runtime() if sim is not None: # Benchmarks must aggregate and report after simulator # teardown; the SimulationManager default exits the whole @@ -459,6 +565,25 @@ def run(self) -> BenchmarkRunResult: self.suite.protocol.measured_trials, ) write_json(run_dir / "aggregates.json", aggregates) + enabled_scenarios = {track.scenario for track in enabled_tracks} + track_notes: list[str] = [] + if "free_space" in enabled_scenarios: + track_notes.append( + "Collision, dynamic, execution, and task metrics are N/A in " + "free-space-common v1." + ) + if "atomic_task" in enabled_scenarios: + track_notes.extend( + [ + "Atomic Task cost_time_ms measures AtomicActionEngine.compile only; " + "execution_time_ms is common physics-replay wall time and " + "end_to_end_time_ms is their sum.", + "trajectory_duration_s is planner-native nominal duration; " + "task_completion_time_s is simulated replay time through the " + "stability hold for successful tasks.", + ] + ) + track_notes.extend(summarize_video_recording(self.video, self._video_paths)) report_path = write_markdown_report( run_dir / "report.md", self.suite, @@ -473,7 +598,7 @@ def run(self) -> BenchmarkRunResult: "path lengths, path_efficiency) average only motion_valid outcomes " "(success-conditioned / survivor-biased). Always read them with n_valid; " "a high path_efficiency on n_valid=2 is not comparable to n_valid=200.", - "Collision, dynamic, execution, and task metrics are N/A in free-space-common v1.", + *track_notes, "Leaderboard and Success-table boolean rates " "(overall_success_rate / success_rate / motion_valid_rate / " "planning_success_rate / ordered_waypoint_success_rate) are macro averages " diff --git a/scripts/benchmark/motion_generation/scenarios/__init__.py b/scripts/benchmark/motion_generation/scenarios/__init__.py index e959d74df..34b4328a3 100644 --- a/scripts/benchmark/motion_generation/scenarios/__init__.py +++ b/scripts/benchmark/motion_generation/scenarios/__init__.py @@ -18,7 +18,33 @@ from __future__ import annotations -from .base import ScenarioProvider +from .atomic_objects import ( + AtomicObjectHandle, + atomic_object_kind_names, + create_atomic_object, + register_atomic_object_kind, +) +from .atomic_task import ( + AtomicSkillCaseProvider, + AtomicTaskScenario, + atomic_skill_provider_names, + create_atomic_skill_provider, + register_atomic_skill_provider, +) +from .base import ScenarioEvaluation, ScenarioProvider from .free_space import FreeSpaceScenario -__all__ = ["FreeSpaceScenario", "ScenarioProvider"] +__all__ = [ + "AtomicObjectHandle", + "AtomicSkillCaseProvider", + "AtomicTaskScenario", + "atomic_object_kind_names", + "atomic_skill_provider_names", + "create_atomic_object", + "create_atomic_skill_provider", + "FreeSpaceScenario", + "register_atomic_object_kind", + "register_atomic_skill_provider", + "ScenarioEvaluation", + "ScenarioProvider", +] diff --git a/scripts/benchmark/motion_generation/scenarios/atomic_objects.py b/scripts/benchmark/motion_generation/scenarios/atomic_objects.py new file mode 100644 index 000000000..cb706a361 --- /dev/null +++ b/scripts/benchmark/motion_generation/scenarios/atomic_objects.py @@ -0,0 +1,194 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Config-driven rigid objects for Atomic Task benchmark cases.""" + +from __future__ import annotations + +import math +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING + +import torch + +from embodichain.data import get_data_path +from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg +from embodichain.lab.sim.shapes import CubeCfg, MeshCfg + +if TYPE_CHECKING: + from embodichain.lab.sim import SimulationManager + from embodichain.lab.sim.objects import RigidObject + +__all__ = [ + "AtomicObjectHandle", + "atomic_object_kind_names", + "create_atomic_object", + "register_atomic_object_kind", +] + +AtomicShapeFactory = Callable[[Mapping[str, object]], object] +_ATOMIC_OBJECT_KINDS: dict[str, AtomicShapeFactory] = {} + + +@dataclass +class AtomicObjectHandle: + """Simulation object plus its algorithm-independent frozen state.""" + + object_id: str + kind: str + config: dict[str, object] + entity: "RigidObject" + initial_pose: torch.Tensor + + def reset(self) -> None: + """Restore the frozen initial pose and clear residual dynamics.""" + self.entity.set_local_pose(self.initial_pose) + self.entity.clear_dynamics() + + def park(self, index: int) -> None: + """Move an inactive object outside every benchmark workspace.""" + pose = self.initial_pose.clone() + pose[:, 0, 3] = 8.0 + float(index) + pose[:, 1, 3] = 8.0 + pose[:, 2, 3] = 1.0 + self.entity.set_local_pose(pose) + self.entity.clear_dynamics() + + +def register_atomic_object_kind(name: str, factory: AtomicShapeFactory) -> None: + """Register a config-only object shape factory.""" + if not name: + raise ValueError("Atomic object kind must not be empty.") + previous = _ATOMIC_OBJECT_KINDS.get(name) + if previous is not None and previous is not factory: + raise ValueError(f"Atomic object kind {name!r} is already registered.") + _ATOMIC_OBJECT_KINDS[name] = factory + + +def atomic_object_kind_names() -> tuple[str, ...]: + """Return registered object kinds in deterministic order.""" + return tuple(sorted(_ATOMIC_OBJECT_KINDS)) + + +def _vector( + value: object, *, name: str, length: int, default: Sequence[float] +) -> list[float]: + """Validate and normalize a numeric vector from YAML configuration.""" + resolved = default if value is None else value + if not isinstance(resolved, Sequence) or isinstance(resolved, (str, bytes)): + raise TypeError(f"{name} must be a sequence of {length} numbers.") + result = [float(item) for item in resolved] + if len(result) != length or not all(math.isfinite(item) for item in result): + raise ValueError(f"{name} must contain {length} finite values.") + return result + + +def _cube_shape(config: Mapping[str, object]) -> CubeCfg: + """Build a cube shape from its declarative size.""" + size = _vector( + config.get("size"), name="cube.size", length=3, default=(0.05, 0.05, 0.05) + ) + if any(value <= 0.0 for value in size): + raise ValueError("cube.size values must be greater than zero.") + return CubeCfg(size=size) + + +def _mesh_shape(config: Mapping[str, object]) -> MeshCfg: + """Build a mesh shape from an absolute or EmbodiChain data path.""" + asset_path = config.get("asset_path") + if not isinstance(asset_path, str) or not asset_path: + raise ValueError("mesh.asset_path must be a non-empty string.") + resolved = Path(asset_path) + return MeshCfg( + fpath=str(resolved if resolved.is_absolute() else get_data_path(asset_path)) + ) + + +def create_atomic_object( + simulation: "SimulationManager", object_config: Mapping[str, object] +) -> AtomicObjectHandle: + """Create one config-driven object and freeze its settled initial pose.""" + object_id = object_config.get("id") + kind = object_config.get("kind") + if not isinstance(object_id, str) or not object_id: + raise ValueError("Every atomic object must define a non-empty id.") + if not isinstance(kind, str) or not kind: + raise ValueError(f"Atomic object {object_id!r} must define a kind.") + try: + factory = _ATOMIC_OBJECT_KINDS[kind] + except KeyError as exc: + raise ValueError( + f"Unknown atomic object kind {kind!r}; registered kinds: " + f"{atomic_object_kind_names()}." + ) from exc + + config = dict(object_config) + position = _vector( + config.get("position"), + name=f"objects[{object_id}].position", + length=3, + default=(-0.42, -0.08, 0.05), + ) + rotation = _vector( + config.get("rotation_deg"), + name=f"objects[{object_id}].rotation_deg", + length=3, + default=(0.0, 0.0, 0.0), + ) + scale = _vector( + config.get("scale"), + name=f"objects[{object_id}].scale", + length=3, + default=(1.0, 1.0, 1.0), + ) + entity = simulation.add_rigid_object( + cfg=RigidObjectCfg( + uid=f"atomic_benchmark_{object_id}", + shape=factory(config), + attrs=RigidBodyAttributesCfg( + mass=float(config.get("mass", 0.05)), + dynamic_friction=float(config.get("dynamic_friction", 0.97)), + static_friction=float(config.get("static_friction", 0.99)), + restitution=float(config.get("restitution", 0.0)), + contact_offset=float(config.get("contact_offset", 0.003)), + rest_offset=float(config.get("rest_offset", 0.001)), + linear_damping=float(config.get("linear_damping", 0.7)), + angular_damping=float(config.get("angular_damping", 0.7)), + min_position_iters=int(config.get("min_position_iters", 32)), + min_velocity_iters=int(config.get("min_velocity_iters", 8)), + ), + max_convex_hull_num=int(config.get("max_convex_hull_num", 16)), + init_pos=position, + init_rot=rotation, + body_scale=scale, + use_usd_properties=bool(config.get("use_usd_properties", False)), + ) + ) + simulation.update(step=int(config.get("settle_steps", 10))) + entity.clear_dynamics() + return AtomicObjectHandle( + object_id=object_id, + kind=kind, + config=config, + entity=entity, + initial_pose=entity.get_local_pose(to_matrix=True).clone(), + ) + + +register_atomic_object_kind("cube", _cube_shape) +register_atomic_object_kind("mesh", _mesh_shape) diff --git a/scripts/benchmark/motion_generation/scenarios/atomic_task.py b/scripts/benchmark/motion_generation/scenarios/atomic_task.py new file mode 100644 index 000000000..6ce1d9bce --- /dev/null +++ b/scripts/benchmark/motion_generation/scenarios/atomic_task.py @@ -0,0 +1,2328 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Physics-backed Atomic Task track shared by planner adapters.""" + +from __future__ import annotations + +import math +import time +from abc import ABC, abstractmethod +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, replace +from pathlib import Path +from typing import TYPE_CHECKING + +import torch + +from embodichain.lab.sim.atomic_actions import ( + ActionInvocation, + Affordance, + AtomicActionEngine, + ControlPartCommandProfile, + EndEffectorPoseGoal, + GraspGoal, + HeldObjectPoseGoal, + HeldObjectState, + JointPositionGoal, + MotionPolicy, + MoveHeldObjectOptions, + MoveJointsOptions, + ObjectSemantics, + PickUpOptions, + PlaceGoal, + PlaceOptions, + PressAffordance, + PressGoal, + PressOptions, + TaskState, + create_simulation_atomic_action_engine, +) +from embodichain.lab.sim.atomic_actions.plans import CompiledTrajectory +from embodichain.lab.sim.planners.utils import PlanResult + +from ..config import SuiteCfg, TrackCfg +from ..metrics.trajectory import compute_case_outcomes +from ..models import BenchmarkCase, CaseOutcome +from ..registry import register_scenario_provider +from ..video import VideoRecordCfg, build_video_path, record_with_window +from .atomic_objects import AtomicObjectHandle, create_atomic_object +from .base import ScenarioEvaluation, ScenarioProvider + +if TYPE_CHECKING: + from embodichain.lab.sim import SimulationManager + from embodichain.lab.sim.objects import Robot + + from ..planners.base import PlannerAdapter + +__all__ = [ + "AtomicSkillCaseProvider", + "AtomicTaskScenario", + "atomic_skill_provider_names", + "create_atomic_skill_provider", + "register_atomic_skill_provider", +] + +_TOP_DOWN_ROTATION = ( + (-0.0539, -0.9985, -0.0022), + (-0.9977, 0.0540, -0.0401), + (0.0401, 0.0000, -0.9992), +) +_TASK_DIFFICULTIES = {"simple", "medium", "hard"} +_GRIPPER_SKILLS = {"pick_up", "move_held_object", "place", "press"} +_CASE_QPOS_RESOLUTION_RAD = 1.0e-4 + + +@dataclass(frozen=True) +class _ExecutionObservation: + """Common-execution measurements used by skill-specific task rules.""" + + execution_success: torch.Tensor + final_tcp_pose: torch.Tensor + joint_tracking_rmse_rad: torch.Tensor + execution_time_ms: float + task_completion_time_s: float + object_lift_delta_m: torch.Tensor | None = None + final_arm_qpos: torch.Tensor | None = None + final_object_pose: torch.Tensor | None = None + + +@dataclass(frozen=True) +class _PhysicsReplaySettings: + """Shared physics-step counts for timed evaluation and video replay.""" + + steps_per_waypoint: int + hold_steps: int + hold_sim_steps: int + joint_tracking_tolerance_rad: float + + +class AtomicSkillCaseProvider(ABC): + """Generate and ground one Atomic Action without planner-specific logic.""" + + skill_id: str + requires_gripper = False + + @abstractmethod + def generate_case( + self, + scenario: "AtomicTaskScenario", + suite: SuiteCfg, + track: TrackCfg, + config: Mapping[str, object], + *, + seed: int, + batch_size: int, + ) -> BenchmarkCase: + """Generate one frozen case and independent IK validity evidence.""" + + @abstractmethod + def build_invocation( + self, + scenario: "AtomicTaskScenario", + case: BenchmarkCase, + adapter: "PlannerAdapter", + ) -> ActionInvocation: + """Ground the case into one planner-independent action invocation.""" + + def object_id(self, case: BenchmarkCase) -> str | None: + """Return the manipulated object identifier when one exists.""" + return case.object_id + + def lift_segment_start(self, compiled: CompiledTrajectory) -> int | None: + """Return the first lift waypoint that should release object dynamics.""" + return None + + def initial_task_state( + self, scenario: "AtomicTaskScenario", case: BenchmarkCase + ) -> TaskState | None: + """Return an optional symbolic precondition for isolated action testing.""" + del scenario, case + return None + + @abstractmethod + def task_result( + self, + scenario: "AtomicTaskScenario", + case: BenchmarkCase, + compiled: CompiledTrajectory, + observation: _ExecutionObservation, + motion_outcomes: tuple[CaseOutcome, ...], + ) -> tuple[torch.Tensor, str]: + """Return per-environment task success and its stable failure code.""" + + +AtomicSkillProviderType = type[AtomicSkillCaseProvider] +_ATOMIC_SKILL_PROVIDERS: dict[str, AtomicSkillProviderType] = {} + + +def register_atomic_skill_provider( + skill_id: str, provider_type: AtomicSkillProviderType +) -> None: + """Register one Atomic Action case provider.""" + if not skill_id: + raise ValueError("Atomic skill id must not be empty.") + previous = _ATOMIC_SKILL_PROVIDERS.get(skill_id) + if previous is not None and previous is not provider_type: + raise ValueError(f"Atomic skill provider {skill_id!r} is already registered.") + _ATOMIC_SKILL_PROVIDERS[skill_id] = provider_type + + +def atomic_skill_provider_names() -> tuple[str, ...]: + """Return registered Atomic Action case providers.""" + return tuple(sorted(_ATOMIC_SKILL_PROVIDERS)) + + +def create_atomic_skill_provider(skill_id: str) -> AtomicSkillCaseProvider: + """Construct one registered Atomic Action case provider.""" + try: + provider_type = _ATOMIC_SKILL_PROVIDERS[skill_id] + except KeyError as exc: + raise ValueError( + f"Unknown atomic skill {skill_id!r}; registered skills: " + f"{atomic_skill_provider_names()}." + ) from exc + return provider_type() + + +def _float_vector( + value: object, + *, + name: str, + length: int, + default: Sequence[float] | None = None, +) -> list[float]: + """Resolve a finite numeric vector from a YAML-compatible value.""" + resolved = default if value is None else value + if not isinstance(resolved, Sequence) or isinstance(resolved, (str, bytes)): + raise TypeError(f"{name} must be a sequence of {length} numbers.") + result = [float(item) for item in resolved] + if len(result) != length or not all(math.isfinite(item) for item in result): + raise ValueError(f"{name} must contain {length} finite values.") + return result + + +def _case_name(config: Mapping[str, object]) -> str: + """Return and validate a stable case name.""" + name = config.get("name") + if not isinstance(name, str) or not name: + raise ValueError("Every atomic skill case must define a non-empty name.") + return name + + +def _difficulty(config: Mapping[str, object]) -> str: + """Resolve the explicit, frozen Atomic Task difficulty label.""" + difficulty = str(config.get("task_difficulty", "simple")) + if difficulty not in _TASK_DIFFICULTIES: + raise ValueError( + f"task_difficulty must be one of {sorted(_TASK_DIFFICULTIES)}." + ) + return difficulty + + +def _seeded_jitter( + amplitude: Sequence[float], + *, + seed: int, + stream: int, + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor: + """Sample deterministic independent uniform jitter in ``[-amplitude, +amplitude]``.""" + generator = torch.Generator(device="cpu") + generator.manual_seed((int(seed) * 1_000_003 + int(stream) * 97_409) % (2**63 - 1)) + unit = torch.rand(len(amplitude), generator=generator, dtype=torch.float32) + values = (2.0 * unit - 1.0) * torch.tensor(amplitude, dtype=torch.float32) + return values.to(dtype=dtype, device=device) + + +def _case_generation_seed(seed: int, *, skill_index: int, case_index: int) -> int: + """Derive one stable RNG seed for all stochastic IK work in a case.""" + return ( + int(seed) * 1_000_003 + int(skill_index) * 97_409 + int(case_index) * 13_007 + ) % (2**63 - 1) + + +def _canonical_case_qpos(qpos: torch.Tensor) -> torch.Tensor: + """Remove insignificant device-level IK noise from frozen case states.""" + return torch.round(qpos / _CASE_QPOS_RESOLUTION_RAD) * _CASE_QPOS_RESOLUTION_RAD + + +def _randomized_vector( + base: Sequence[float], + config: Mapping[str, object], + *, + jitter_name: str, + seed: int, + stream: int, + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor: + """Return one configured vector plus optional deterministic uniform jitter.""" + base_values = [float(value) for value in base] + amplitude = _float_vector( + config.get(jitter_name), + name=jitter_name, + length=len(base_values), + default=(0.0,) * len(base_values), + ) + if any(value < 0.0 for value in amplitude): + raise ValueError(f"{jitter_name} values must be non-negative.") + return torch.tensor(base_values, dtype=dtype, device=device) + _seeded_jitter( + amplitude, + seed=seed, + stream=stream, + dtype=dtype, + device=device, + ) + + +def _randomization_parameters( + config: Mapping[str, object], *, seed: int +) -> dict[str, object]: + """Serialize the deterministic randomization contract into a case manifest.""" + ranges = { + str(key): value + for key, value in config.items() + if str(key).endswith("_jitter_m") or str(key).endswith("_jitter_rad") + } + return { + "enabled": bool(ranges), + "seed": int(seed), + "distribution": "independent_uniform", + "ranges": ranges, + } + + +def _motion_valid_mask( + motion_outcomes: tuple[CaseOutcome, ...], *, device: torch.device +) -> torch.Tensor: + """Return the common motion-valid result as one device-local mask.""" + return torch.tensor( + [item.motion_valid for item in motion_outcomes], + dtype=torch.bool, + device=device, + ) + + +def _case_pose(case: BenchmarkCase, name: str, *, device: torch.device) -> torch.Tensor: + """Restore one frozen pose tensor from JSON-compatible case parameters.""" + return torch.tensor(case.case_parameters[name], dtype=torch.float32, device=device) + + +def _held_object_state( + scenario: "AtomicTaskScenario", case: BenchmarkCase +) -> HeldObjectState: + """Reconstruct the frozen held-object relation for an isolated case.""" + handle = scenario.object_handle(case.object_id) + semantics = ObjectSemantics( + affordance=Affordance(), + geometry={}, + properties={"benchmark_object_id": handle.object_id}, + label=handle.object_id, + entity_id=handle.entity.uid, + ) + return HeldObjectState( + semantics=semantics, + object_to_eef=_case_pose(case, "object_to_eef", device=scenario.robot.device), + grasp_xpos=_case_pose(case, "grasp_pose", device=scenario.robot.device), + ) + + +class _MoveEndEffectorCases(AtomicSkillCaseProvider): + """Deterministic robot-relative MoveEndEffector cases.""" + + skill_id = "move_end_effector" + + def generate_case( + self, + scenario: "AtomicTaskScenario", + suite: SuiteCfg, + track: TrackCfg, + config: Mapping[str, object], + *, + seed: int, + batch_size: int, + ) -> BenchmarkCase: + scenario.restore_base_robot() + scenario.randomize_robot_start(config, seed=seed, stream=11) + raw_offsets = config.get("target_offsets_m") + if not isinstance(raw_offsets, Sequence) or isinstance( + raw_offsets, (str, bytes) + ): + raise TypeError("target_offsets_m must be a non-empty list of xyz vectors.") + base_offsets = [ + _float_vector(value, name="target_offsets_m", length=3) + for value in raw_offsets + ] + if not base_offsets: + raise ValueError("target_offsets_m must not be empty.") + + start_qpos = scenario.robot.get_qpos(name=scenario.control_part).clone() + start_pose = scenario.robot.compute_fk( + start_qpos, name=scenario.control_part, to_matrix=True + ) + offsets_tensor = torch.stack( + [ + _randomized_vector( + offset, + config, + jitter_name="target_offset_jitter_m", + seed=seed, + stream=101 + index, + dtype=start_pose.dtype, + device=start_pose.device, + ) + for index, offset in enumerate(base_offsets) + ] + ) + offsets = offsets_tensor.detach().cpu().tolist() + targets = start_pose[:, None].repeat(1, len(offsets), 1, 1) + targets[:, :, :3, 3] += offsets_tensor[None] + references = scenario.solve_reference_qpos(start_qpos, targets) + name = _case_name(config) + return BenchmarkCase( + suite_version=suite.suite_version, + track=track.id, + scenario_id=self.skill_id, + case_id=f"{track.id}:{self.skill_id}:{name}:s{seed}", + seed=seed, + batch_size=batch_size, + num_waypoints=len(offsets), + path_shape="robot_relative_waypoints", + start_state_bin="pre_action", + start_qpos=start_qpos, + target_waypoints=targets, + reference_qpos=references, + robot_id=suite.robot.id, + skill_id=self.skill_id, + task_difficulty=_difficulty(config), + primary_success="task_success", + full_start_qpos=_canonical_case_qpos(scenario.robot.get_qpos().clone()), + case_parameters={ + "sample_count": int(config.get("sample_count", 80)), + "target_offsets_m": offsets, + "randomization": _randomization_parameters(config, seed=seed), + "difficulty_factors": dict(config.get("difficulty_factors", {})), + }, + ) + + def build_invocation( + self, + scenario: "AtomicTaskScenario", + case: BenchmarkCase, + adapter: "PlannerAdapter", + ) -> ActionInvocation: + return scenario.require_engine().make_invocation( + self.skill_id, + EndEffectorPoseGoal(case.target_waypoints), + control_parts={"primary": {"motion": scenario.control_part}}, + motion_policy=MotionPolicy( + strategy="motion_gen", + sample_count=int(case.case_parameters["sample_count"]), + ), + ) + + def task_result( + self, + scenario: "AtomicTaskScenario", + case: BenchmarkCase, + compiled: CompiledTrajectory, + observation: _ExecutionObservation, + motion_outcomes: tuple[CaseOutcome, ...], + ) -> tuple[torch.Tensor, str]: + del compiled + target = case.target_waypoints[:, -1] + translation = torch.linalg.vector_norm( + observation.final_tcp_pose[:, :3, 3] - target[:, :3, 3], dim=-1 + ) + relative = ( + target[:, :3, :3].transpose(-1, -2) @ observation.final_tcp_pose[:, :3, :3] + ) + trace = torch.diagonal(relative, dim1=-2, dim2=-1).sum(dim=-1) + rotation = torch.arccos(torch.clamp((trace - 1.0) * 0.5, -1.0, 1.0)) + motion_valid = torch.tensor( + [item.motion_valid for item in motion_outcomes], + dtype=torch.bool, + device=translation.device, + ) + success = ( + observation.execution_success + & motion_valid + & (translation <= scenario.suite.protocol.position_threshold_m) + & (rotation <= scenario.suite.protocol.rotation_threshold_rad) + ) + return success, "task_goal_miss" + + +class _PickUpCases(AtomicSkillCaseProvider): + """Explicit-grasp PickUp cases that isolate motion-planner performance.""" + + skill_id = "pick_up" + requires_gripper = True + + def generate_case( + self, + scenario: "AtomicTaskScenario", + suite: SuiteCfg, + track: TrackCfg, + config: Mapping[str, object], + *, + seed: int, + batch_size: int, + ) -> BenchmarkCase: + object_id = config.get("object") + if not isinstance(object_id, str) or not object_id: + raise ValueError("PickUp cases must reference a non-empty object id.") + handle = scenario.activate_object(object_id) + scenario.restore_base_robot() + + object_pose = scenario.randomize_object_pose( + handle, config, seed=seed, stream=201 + ) + arm_start = scenario.robot.get_qpos(name=scenario.control_part) + pre_pick_pose = scenario.robot.compute_fk( + arm_start, name=scenario.control_part, to_matrix=True + ).clone() + pre_pick_pose[:, :2, 3] = object_pose[:, :2, 3] + pre_pick_pose[:, 2, 3] = float(config.get("pre_pick_height_m", 0.36)) + success, pre_pick_qpos = scenario.robot.compute_ik( + pose=pre_pick_pose, + joint_seed=arm_start, + name=scenario.control_part, + ) + if not bool(torch.as_tensor(success).all().item()): + raise RuntimeError( + f"Independent IK rejected PickUp case {_case_name(config)!r}." + ) + pre_pick_qpos = _canonical_case_qpos(pre_pick_qpos) + scenario.set_robot_start(pre_pick_qpos, open_gripper=True) + + approach = torch.tensor( + _float_vector( + config.get("approach_direction"), + name="approach_direction", + length=3, + default=(0.0, 0.0, -1.0), + ), + dtype=object_pose.dtype, + device=object_pose.device, + ) + approach_norm = torch.linalg.vector_norm(approach) + if float(approach_norm.item()) <= 1.0e-6: + raise ValueError("approach_direction must be non-zero.") + approach = approach / approach_norm + pre_grasp_distance = float(config.get("pre_grasp_distance_m", 0.15)) + lift_height = float(config.get("lift_height_m", 0.16)) + + grasp_source = str(config.get("grasp_source", "fixed")) + if grasp_source == "antipodal": + grasp_pose = scenario.resolve_antipodal_grasp( + handle, + object_pose, + approach, + seed=seed, + start_qpos=scenario.robot.get_qpos(name=scenario.control_part), + pre_grasp_distance=pre_grasp_distance, + lift_height=lift_height, + n_sample=int(config.get("grasp_sample_count", 10_000)), + max_candidates=int(config.get("grasp_max_candidates", 128)), + alignment_max_angle_deg=float( + config.get("grasp_alignment_max_angle_deg", 10.0) + ), + ) + elif grasp_source == "fixed": + grasp_pose = object_pose.clone() + rotation_value = config.get("grasp_rotation", _TOP_DOWN_ROTATION) + if not isinstance(rotation_value, Sequence) or len(rotation_value) != 3: + raise ValueError("grasp_rotation must be a 3x3 matrix.") + rotation = torch.tensor( + rotation_value, dtype=grasp_pose.dtype, device=grasp_pose.device + ) + if rotation.shape != (3, 3): + raise ValueError("grasp_rotation must be a 3x3 matrix.") + grasp_pose[:, :3, :3] = rotation + else: + raise ValueError("grasp_source must be 'fixed' or 'antipodal'.") + grasp_offset = _randomized_vector( + _float_vector( + config.get("grasp_offset_m"), + name="grasp_offset_m", + length=3, + default=(0.0, 0.0, 0.0), + ), + config, + jitter_name="grasp_offset_jitter_m", + seed=seed, + stream=202, + dtype=grasp_pose.dtype, + device=grasp_pose.device, + ) + grasp_pose[:, :3, 3] += grasp_offset + pre_grasp = grasp_pose.clone() + pre_grasp[:, :3, 3] -= approach * pre_grasp_distance + lift = grasp_pose.clone() + lift[:, 2, 3] += lift_height + targets = torch.stack([pre_grasp, grasp_pose, lift], dim=1) + start_qpos = scenario.robot.get_qpos(name=scenario.control_part).clone() + references = scenario.solve_reference_qpos(start_qpos, targets) + name = _case_name(config) + return BenchmarkCase( + suite_version=suite.suite_version, + track=track.id, + scenario_id=self.skill_id, + case_id=f"{track.id}:{self.skill_id}:{name}:s{seed}", + seed=seed, + batch_size=batch_size, + num_waypoints=3, + path_shape="approach_grasp_lift", + start_state_bin="pre_pick", + start_qpos=start_qpos, + target_waypoints=targets, + reference_qpos=references, + robot_id=suite.robot.id, + skill_id=self.skill_id, + object_id=object_id, + task_difficulty=_difficulty(config), + primary_success="task_success", + full_start_qpos=_canonical_case_qpos(scenario.robot.get_qpos().clone()), + case_parameters={ + "sample_count": int(config.get("sample_count", 120)), + "grasp_source": grasp_source, + "hand_interp_steps": int(config.get("hand_interp_steps", 12)), + "approach_direction": approach.detach().cpu().tolist(), + "pre_grasp_distance_m": pre_grasp_distance, + "lift_height_m": lift_height, + "minimum_object_lift_m": float( + config.get("minimum_object_lift_m", 0.04) + ), + "grasp_pose": grasp_pose.detach().cpu().tolist(), + "object_initial_pose": object_pose.detach().cpu().tolist(), + "object_config": dict(handle.config), + "randomization": _randomization_parameters(config, seed=seed), + "difficulty_factors": dict(config.get("difficulty_factors", {})), + }, + ) + + def build_invocation( + self, + scenario: "AtomicTaskScenario", + case: BenchmarkCase, + adapter: "PlannerAdapter", + ) -> ActionInvocation: + handle = scenario.object_handle(case.object_id) + if scenario.end_effector_part is None: + raise RuntimeError( + "PickUp requires a configured end-effector control part." + ) + semantics = ObjectSemantics( + affordance=Affordance(), + geometry={}, + properties={"benchmark_object_id": handle.object_id}, + label=handle.object_id, + entity_id=handle.entity.uid, + ) + return scenario.require_engine().make_invocation( + self.skill_id, + GraspGoal( + semantics=semantics, + grasp_xpos=case.target_waypoints[:, 1], + ), + control_parts={ + "primary": { + "motion": scenario.control_part, + "grasp": scenario.end_effector_part, + } + }, + motion_policy=MotionPolicy( + strategy="motion_gen", + sample_count=int(case.case_parameters["sample_count"]), + ), + skill_options=PickUpOptions( + approach_direction=torch.tensor( + case.case_parameters["approach_direction"], + dtype=torch.float32, + device=scenario.robot.device, + ), + pre_grasp_distance=float(case.case_parameters["pre_grasp_distance_m"]), + lift_height=float(case.case_parameters["lift_height_m"]), + hand_interp_steps=int(case.case_parameters["hand_interp_steps"]), + ), + ) + + def lift_segment_start(self, compiled: CompiledTrajectory) -> int | None: + """Return the lift boundary emitted by PickUp.""" + return compiled.segment(0, "lift").start + + def task_result( + self, + scenario: "AtomicTaskScenario", + case: BenchmarkCase, + compiled: CompiledTrajectory, + observation: _ExecutionObservation, + motion_outcomes: tuple[CaseOutcome, ...], + ) -> tuple[torch.Tensor, str]: + held_created = ( + compiled.projected_context.get_held_object(scenario.control_part) + is not None + ) + lift = observation.object_lift_delta_m + if lift is None: + lift = torch.full( + (case.batch_size,), + -torch.inf, + device=observation.execution_success.device, + ) + motion_valid = torch.tensor( + [item.motion_valid for item in motion_outcomes], + dtype=torch.bool, + device=lift.device, + ) + success = ( + observation.execution_success + & motion_valid + & held_created + & (lift >= float(case.case_parameters["minimum_object_lift_m"])) + ) + return success, "object_not_grasped" + + +class _MoveJointsCases(AtomicSkillCaseProvider): + """Deterministic relative joint-space waypoint cases.""" + + skill_id = "move_joints" + + def generate_case( + self, + scenario: "AtomicTaskScenario", + suite: SuiteCfg, + track: TrackCfg, + config: Mapping[str, object], + *, + seed: int, + batch_size: int, + ) -> BenchmarkCase: + scenario.restore_base_robot() + scenario.randomize_robot_start(config, seed=seed, stream=21) + raw_offsets = config.get("target_offsets_rad") + if not isinstance(raw_offsets, Sequence) or isinstance( + raw_offsets, (str, bytes) + ): + raise TypeError("target_offsets_rad must be a non-empty list.") + base_offsets = [ + _float_vector(value, name="target_offsets_rad", length=7) + for value in raw_offsets + ] + if not base_offsets: + raise ValueError("target_offsets_rad must not be empty.") + start_qpos = scenario.robot.get_qpos(name=scenario.control_part).clone() + offset_tensor = torch.stack( + [ + _randomized_vector( + offset, + config, + jitter_name="target_offset_jitter_rad", + seed=seed, + stream=211 + index, + dtype=start_qpos.dtype, + device=start_qpos.device, + ) + for index, offset in enumerate(base_offsets) + ] + ) + offsets = offset_tensor.detach().cpu().tolist() + targets = start_qpos[:, None] + torch.cumsum(offset_tensor, dim=0)[None] + limits = scenario.robot.get_qpos_limits(name=scenario.control_part)[0] + margin = float(config.get("joint_limit_margin_rad", 0.05)) + if bool( + ( + (targets < limits[:, 0][None, None] + margin) + | (targets > limits[:, 1][None, None] - margin) + ) + .any() + .item() + ): + raise RuntimeError( + f"Joint limits rejected MoveJoints case {_case_name(config)!r}." + ) + target_poses = torch.stack( + [ + scenario.robot.compute_fk( + targets[:, index], name=scenario.control_part, to_matrix=True + ) + for index in range(targets.shape[1]) + ], + dim=1, + ) + name = _case_name(config) + return BenchmarkCase( + suite_version=suite.suite_version, + track=track.id, + scenario_id=self.skill_id, + case_id=f"{track.id}:{self.skill_id}:{name}:s{seed}", + seed=seed, + batch_size=batch_size, + num_waypoints=len(offsets), + path_shape="relative_joint_waypoints", + start_state_bin="pre_action", + start_qpos=start_qpos, + target_waypoints=target_poses, + reference_qpos=targets, + robot_id=suite.robot.id, + skill_id=self.skill_id, + task_difficulty=_difficulty(config), + primary_success="task_success", + full_start_qpos=_canonical_case_qpos(scenario.robot.get_qpos().clone()), + case_parameters={ + "sample_count": int(config.get("sample_count", 80)), + "target_offsets_rad": offsets, + "target_qpos": targets.detach().cpu().tolist(), + "joint_threshold_rad": float(config.get("joint_threshold_rad", 0.02)), + "randomization": _randomization_parameters(config, seed=seed), + "difficulty_factors": dict(config.get("difficulty_factors", {})), + }, + ) + + def build_invocation( + self, + scenario: "AtomicTaskScenario", + case: BenchmarkCase, + adapter: "PlannerAdapter", + ) -> ActionInvocation: + target = torch.tensor( + case.case_parameters["target_qpos"], + dtype=torch.float32, + device=scenario.robot.device, + ) + return scenario.require_engine().make_invocation( + self.skill_id, + JointPositionGoal(target), + control_parts={"primary": {"motion": scenario.control_part}}, + motion_policy=MotionPolicy( + strategy="motion_gen", + sample_count=int(case.case_parameters["sample_count"]), + ), + skill_options=MoveJointsOptions(), + ) + + def task_result( + self, + scenario: "AtomicTaskScenario", + case: BenchmarkCase, + compiled: CompiledTrajectory, + observation: _ExecutionObservation, + motion_outcomes: tuple[CaseOutcome, ...], + ) -> tuple[torch.Tensor, str]: + del scenario, compiled + if observation.final_arm_qpos is None: + return torch.zeros_like(observation.execution_success), "joint_goal_miss" + target = torch.tensor( + case.case_parameters["target_qpos"], + dtype=observation.final_arm_qpos.dtype, + device=observation.final_arm_qpos.device, + )[:, -1] + error = torch.amax(torch.abs(observation.final_arm_qpos - target), dim=-1) + success = ( + observation.execution_success + & _motion_valid_mask(motion_outcomes, device=error.device) + & (error <= float(case.case_parameters["joint_threshold_rad"])) + ) + return success, "joint_goal_miss" + + +class _HeldObjectCases(AtomicSkillCaseProvider): + """Shared isolated held-object precondition for transport and placement.""" + + requires_gripper = True + + def _prepare_start( + self, + scenario: "AtomicTaskScenario", + config: Mapping[str, object], + *, + seed: int, + ) -> tuple[ + AtomicObjectHandle, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + ]: + object_id = config.get("object") + if not isinstance(object_id, str) or not object_id: + raise ValueError(f"{self.skill_id} cases require an object id.") + handle = scenario.activate_object(object_id) + scenario.restore_base_robot() + table_object_pose = scenario.randomize_object_pose( + handle, config, seed=seed, stream=301 + ) + object_pose = table_object_pose.clone() + held_offset = _randomized_vector( + _float_vector( + config.get("held_object_offset_m"), + name="held_object_offset_m", + length=3, + default=(0.0, 0.0, 0.18), + ), + config, + jitter_name="held_object_offset_jitter_m", + seed=seed, + stream=302, + dtype=object_pose.dtype, + device=object_pose.device, + ) + object_pose[:, :3, 3] += held_offset + grasp_pose = object_pose.clone() + grasp_pose[:, :3, :3] = torch.tensor( + config.get("grasp_rotation", _TOP_DOWN_ROTATION), + dtype=grasp_pose.dtype, + device=grasp_pose.device, + ) + grasp_pose[:, :3, 3] += _randomized_vector( + _float_vector( + config.get("grasp_offset_m"), + name="grasp_offset_m", + length=3, + default=(0.0, 0.0, 0.0), + ), + config, + jitter_name="grasp_offset_jitter_m", + seed=seed, + stream=303, + dtype=grasp_pose.dtype, + device=grasp_pose.device, + ) + arm_seed = scenario.robot.get_qpos(name=scenario.control_part) + success, arm_start = scenario.robot.compute_ik( + pose=grasp_pose, + joint_seed=arm_seed, + name=scenario.control_part, + ) + if not bool(torch.as_tensor(success).all().item()): + raise RuntimeError( + f"Independent IK rejected held-object start {_case_name(config)!r}." + ) + arm_start = _canonical_case_qpos(arm_start) + scenario.set_robot_start(arm_start, open_gripper=False, grasp_gripper=True) + handle.entity.set_local_pose(object_pose) + handle.entity.clear_dynamics() + if scenario.simulation is not None: + scenario.simulation.update(step=2) + handle.entity.set_local_pose(object_pose) + handle.entity.clear_dynamics() + scenario.set_robot_start( + arm_start, + open_gripper=False, + grasp_gripper=True, + ) + object_to_eef = torch.bmm(torch.linalg.inv(object_pose), grasp_pose) + return handle, object_pose, grasp_pose, object_to_eef, table_object_pose + + def initial_task_state( + self, scenario: "AtomicTaskScenario", case: BenchmarkCase + ) -> TaskState: + return TaskState( + batch_size=case.batch_size, + device=scenario.robot.device, + held_objects={scenario.control_part: _held_object_state(scenario, case)}, + ) + + @staticmethod + def _held_parameters( + handle: AtomicObjectHandle, + object_pose: torch.Tensor, + grasp_pose: torch.Tensor, + object_to_eef: torch.Tensor, + ) -> dict[str, object]: + return { + "object_initial_pose": object_pose.detach().cpu().tolist(), + "grasp_pose": grasp_pose.detach().cpu().tolist(), + "object_to_eef": object_to_eef.detach().cpu().tolist(), + "object_config": dict(handle.config), + } + + +class _MoveHeldObjectCases(_HeldObjectCases): + """Move an already held rigid object to one frozen target pose.""" + + skill_id = "move_held_object" + + def generate_case( + self, + scenario: "AtomicTaskScenario", + suite: SuiteCfg, + track: TrackCfg, + config: Mapping[str, object], + *, + seed: int, + batch_size: int, + ) -> BenchmarkCase: + ( + handle, + object_pose, + grasp_pose, + object_to_eef, + _table_object_pose, + ) = self._prepare_start(scenario, config, seed=seed) + target_object = object_pose.clone() + target_object[:, :3, 3] += _randomized_vector( + _float_vector( + config.get("target_object_offset_m"), + name="target_object_offset_m", + length=3, + default=(0.08, 0.08, 0.04), + ), + config, + jitter_name="target_object_offset_jitter_m", + seed=seed, + stream=311, + dtype=object_pose.dtype, + device=object_pose.device, + ) + target_eef = torch.bmm(target_object, object_to_eef) + start_qpos = scenario.robot.get_qpos(name=scenario.control_part).clone() + references = scenario.solve_reference_qpos(start_qpos, target_eef[:, None]) + name = _case_name(config) + return BenchmarkCase( + suite_version=suite.suite_version, + track=track.id, + scenario_id=self.skill_id, + case_id=f"{track.id}:{self.skill_id}:{name}:s{seed}", + seed=seed, + batch_size=batch_size, + num_waypoints=1, + path_shape="held_object_transport", + start_state_bin="object_held", + start_qpos=start_qpos, + target_waypoints=target_eef[:, None], + reference_qpos=references, + robot_id=suite.robot.id, + skill_id=self.skill_id, + object_id=handle.object_id, + task_difficulty=_difficulty(config), + primary_success="task_success", + full_start_qpos=_canonical_case_qpos(scenario.robot.get_qpos().clone()), + case_parameters={ + **self._held_parameters(handle, object_pose, grasp_pose, object_to_eef), + "sample_count": int(config.get("sample_count", 80)), + "target_object_pose": target_object.detach().cpu().tolist(), + "object_position_threshold_m": float( + config.get("object_position_threshold_m", 0.04) + ), + "randomization": _randomization_parameters(config, seed=seed), + "difficulty_factors": dict(config.get("difficulty_factors", {})), + }, + ) + + def build_invocation( + self, + scenario: "AtomicTaskScenario", + case: BenchmarkCase, + adapter: "PlannerAdapter", + ) -> ActionInvocation: + return scenario.require_engine().make_invocation( + self.skill_id, + HeldObjectPoseGoal( + _case_pose(case, "target_object_pose", device=scenario.robot.device) + ), + control_parts={ + "primary": { + "motion": scenario.control_part, + "grasp": scenario.end_effector_part, + } + }, + motion_policy=MotionPolicy( + strategy="motion_gen", + sample_count=int(case.case_parameters["sample_count"]), + ), + skill_options=MoveHeldObjectOptions(), + ) + + def task_result( + self, + scenario: "AtomicTaskScenario", + case: BenchmarkCase, + compiled: CompiledTrajectory, + observation: _ExecutionObservation, + motion_outcomes: tuple[CaseOutcome, ...], + ) -> tuple[torch.Tensor, str]: + final_object = observation.final_object_pose + if final_object is None: + return torch.zeros_like(observation.execution_success), "object_goal_miss" + target = _case_pose(case, "target_object_pose", device=final_object.device) + error = torch.linalg.vector_norm( + final_object[:, :3, 3] - target[:, :3, 3], dim=-1 + ) + held = compiled.projected_context.get_held_object(scenario.control_part) + success = ( + observation.execution_success + & _motion_valid_mask(motion_outcomes, device=error.device) + & (held is not None) + & (error <= float(case.case_parameters["object_position_threshold_m"])) + ) + return success, "object_goal_miss" + + +class _PlaceCases(_HeldObjectCases): + """Place an already held rigid object and release it under physics.""" + + skill_id = "place" + + def generate_case( + self, + scenario: "AtomicTaskScenario", + suite: SuiteCfg, + track: TrackCfg, + config: Mapping[str, object], + *, + seed: int, + batch_size: int, + ) -> BenchmarkCase: + ( + handle, + object_pose, + grasp_pose, + object_to_eef, + table_object_pose, + ) = self._prepare_start(scenario, config, seed=seed) + target_object = table_object_pose.clone() + target_object[:, :3, 3] += _randomized_vector( + _float_vector( + config.get("target_object_offset_m"), + name="target_object_offset_m", + length=3, + default=(0.10, 0.10, 0.0), + ), + config, + jitter_name="target_object_offset_jitter_m", + seed=seed, + stream=321, + dtype=object_pose.dtype, + device=object_pose.device, + ) + release = torch.bmm(target_object, object_to_eef) + lift_height = float(config.get("retract_height_m", 0.10)) + approach = release.clone() + approach[:, 2, 3] += lift_height + retract = approach.clone() + targets = torch.stack([approach, release, retract], dim=1) + start_qpos = scenario.robot.get_qpos(name=scenario.control_part).clone() + references = scenario.solve_reference_qpos(start_qpos, targets) + name = _case_name(config) + return BenchmarkCase( + suite_version=suite.suite_version, + track=track.id, + scenario_id=self.skill_id, + case_id=f"{track.id}:{self.skill_id}:{name}:s{seed}", + seed=seed, + batch_size=batch_size, + num_waypoints=3, + path_shape="approach_release_retract", + start_state_bin="object_held", + start_qpos=start_qpos, + target_waypoints=targets, + reference_qpos=references, + robot_id=suite.robot.id, + skill_id=self.skill_id, + object_id=handle.object_id, + task_difficulty=_difficulty(config), + primary_success="task_success", + full_start_qpos=_canonical_case_qpos(scenario.robot.get_qpos().clone()), + case_parameters={ + **self._held_parameters(handle, object_pose, grasp_pose, object_to_eef), + "sample_count": int(config.get("sample_count", 120)), + "release_pose": release.detach().cpu().tolist(), + "target_object_pose": target_object.detach().cpu().tolist(), + "hand_interp_steps": int(config.get("hand_interp_steps", 12)), + "retract_height_m": lift_height, + "object_position_threshold_m": float( + config.get("object_position_threshold_m", 0.05) + ), + "randomization": _randomization_parameters(config, seed=seed), + "difficulty_factors": dict(config.get("difficulty_factors", {})), + }, + ) + + def build_invocation( + self, + scenario: "AtomicTaskScenario", + case: BenchmarkCase, + adapter: "PlannerAdapter", + ) -> ActionInvocation: + return scenario.require_engine().make_invocation( + self.skill_id, + PlaceGoal(_case_pose(case, "release_pose", device=scenario.robot.device)), + control_parts={ + "primary": { + "motion": scenario.control_part, + "grasp": scenario.end_effector_part, + } + }, + motion_policy=MotionPolicy( + strategy="motion_gen", + sample_count=int(case.case_parameters["sample_count"]), + ), + skill_options=PlaceOptions( + hand_interp_steps=int(case.case_parameters["hand_interp_steps"]), + lift_height=float(case.case_parameters["retract_height_m"]), + ), + ) + + def task_result( + self, + scenario: "AtomicTaskScenario", + case: BenchmarkCase, + compiled: CompiledTrajectory, + observation: _ExecutionObservation, + motion_outcomes: tuple[CaseOutcome, ...], + ) -> tuple[torch.Tensor, str]: + final_object = observation.final_object_pose + if final_object is None: + return torch.zeros_like(observation.execution_success), "object_not_placed" + target = _case_pose(case, "target_object_pose", device=final_object.device) + error = torch.linalg.vector_norm( + final_object[:, :3, 3] - target[:, :3, 3], dim=-1 + ) + released = ( + compiled.projected_context.get_held_object(scenario.control_part) is None + ) + success = ( + observation.execution_success + & _motion_valid_mask(motion_outcomes, device=error.device) + & released + & (error <= float(case.case_parameters["object_position_threshold_m"])) + ) + return success, "object_not_placed" + + +class _PressCases(AtomicSkillCaseProvider): + """Close the gripper, reach a contact pose, and retract.""" + + skill_id = "press" + requires_gripper = True + + def generate_case( + self, + scenario: "AtomicTaskScenario", + suite: SuiteCfg, + track: TrackCfg, + config: Mapping[str, object], + *, + seed: int, + batch_size: int, + ) -> BenchmarkCase: + scenario.restore_base_robot() + scenario.randomize_robot_start(config, seed=seed, stream=31) + start_qpos = scenario.robot.get_qpos(name=scenario.control_part).clone() + start_pose = scenario.robot.compute_fk( + start_qpos, name=scenario.control_part, to_matrix=True + ) + target_pose = start_pose.clone() + target_pose[:, :3, 3] += _randomized_vector( + _float_vector( + config.get("target_offset_m"), + name="target_offset_m", + length=3, + default=(0.0, 0.0, -0.08), + ), + config, + jitter_name="target_offset_jitter_m", + seed=seed, + stream=331, + dtype=start_pose.dtype, + device=start_pose.device, + ) + press_axis = _float_vector( + config.get("press_axis"), + name="press_axis", + length=3, + default=(0.0, 0.0, 1.0), + ) + press_position = _float_vector( + config.get("press_position"), + name="press_position", + length=3, + default=(0.0, 0.0, 0.0), + ) + options = PressOptions( + hand_interp_steps=int(config.get("hand_interp_steps", 8)), + approach_distance=float(config.get("approach_distance_m", 0.1)), + press_distance=float(config.get("press_distance_m", 0.05)), + press_position=tuple(press_position), + ) + affordance = PressAffordance( + press_axis=torch.tensor( + press_axis, + dtype=target_pose.dtype, + device=target_pose.device, + ), + press_position=tuple(press_position), + ) + contact_pose = affordance.get_press_pose( + target_pose, + press_position=options.press_position, + ) + approach_pose = contact_pose.clone() + approach_pose[:, :3, 3] -= contact_pose[:, :3, 2] * options.approach_distance + pressed_pose = contact_pose.clone() + pressed_pose[:, :3, 3] += contact_pose[:, :3, 2] * options.press_distance + targets = torch.stack( + [approach_pose, contact_pose, pressed_pose, approach_pose], + dim=1, + ) + references = scenario.solve_reference_qpos(start_qpos, targets) + name = _case_name(config) + return BenchmarkCase( + suite_version=suite.suite_version, + track=track.id, + scenario_id=self.skill_id, + case_id=f"{track.id}:{self.skill_id}:{name}:s{seed}", + seed=seed, + batch_size=batch_size, + num_waypoints=4, + path_shape="approach_contact_press_retract", + start_state_bin="pre_action", + start_qpos=start_qpos, + target_waypoints=targets, + reference_qpos=references, + robot_id=suite.robot.id, + skill_id=self.skill_id, + task_difficulty=_difficulty(config), + primary_success="task_success", + full_start_qpos=_canonical_case_qpos(scenario.robot.get_qpos().clone()), + case_parameters={ + "sample_count": int(config.get("sample_count", 80)), + "press_target_pose": target_pose.detach().cpu().tolist(), + "press_axis": press_axis, + "press_position": press_position, + "hand_interp_steps": options.hand_interp_steps, + "approach_distance_m": options.approach_distance, + "press_distance_m": options.press_distance, + "waypoint_rotation_symmetry": "half_turn_about_z", + "randomization": _randomization_parameters(config, seed=seed), + "difficulty_factors": dict(config.get("difficulty_factors", {})), + }, + ) + + def build_invocation( + self, + scenario: "AtomicTaskScenario", + case: BenchmarkCase, + adapter: "PlannerAdapter", + ) -> ActionInvocation: + press_axis = torch.tensor( + case.case_parameters["press_axis"], + dtype=torch.float32, + device=scenario.robot.device, + ) + press_position = tuple(case.case_parameters["press_position"]) + semantics = ObjectSemantics( + affordance=PressAffordance( + press_axis=press_axis, + press_position=press_position, + ), + geometry={}, + entity_id=f"benchmark.press_target.{case.case_id}", + label="benchmark_press_target", + ) + return scenario.require_engine().make_invocation( + self.skill_id, + PressGoal( + semantics, + _case_pose( + case, + "press_target_pose", + device=scenario.robot.device, + ), + ), + control_parts={ + "primary": { + "motion": scenario.control_part, + "grasp": scenario.end_effector_part, + } + }, + motion_policy=MotionPolicy( + strategy="motion_gen", + sample_count=int(case.case_parameters["sample_count"]), + ), + skill_options=PressOptions( + hand_interp_steps=int(case.case_parameters["hand_interp_steps"]), + approach_distance=float(case.case_parameters["approach_distance_m"]), + press_distance=float(case.case_parameters["press_distance_m"]), + press_position=press_position, + ), + ) + + def task_result( + self, + scenario: "AtomicTaskScenario", + case: BenchmarkCase, + compiled: CompiledTrajectory, + observation: _ExecutionObservation, + motion_outcomes: tuple[CaseOutcome, ...], + ) -> tuple[torch.Tensor, str]: + del scenario, case, compiled + success = observation.execution_success & _motion_valid_mask( + motion_outcomes, device=observation.execution_success.device + ) + return success, "press_pose_miss" + + +class AtomicTaskScenario(ScenarioProvider): + """Run fixed Atomic Actions through an adapter-owned MotionGenerator.""" + + required_capabilities = frozenset( + {"eef_waypoint", "joint_waypoint", "atomic_action"} + ) + + def __init__(self) -> None: + self.simulation: SimulationManager | None = None + self.robot: Robot | None = None + self.suite: SuiteCfg | None = None + self.track: TrackCfg | None = None + self.control_part = "arm" + self.end_effector_part: str | None = None + self._objects: dict[str, AtomicObjectHandle] = {} + self._case_providers: dict[str, AtomicSkillCaseProvider] = {} + self._base_full_qpos: torch.Tensor | None = None + self._gripper_open: torch.Tensor | None = None + self._gripper_grasp: torch.Tensor | None = None + self._engine: AtomicActionEngine | None = None + + def batch_sizes(self, suite: SuiteCfg, track: TrackCfg) -> list[int]: + """Return the explicitly configured physical-execution batch sizes.""" + del suite + values = [int(value) for value in track.config.get("batch_sizes", [1])] + if values != [1]: + raise ValueError( + "The initial Atomic Task implementation supports batch_sizes: [1] only." + ) + return values + + def configure_runtime( + self, + simulation: SimulationManager, + robot: Robot, + suite: SuiteCfg, + track: TrackCfg, + control_part: str, + ) -> None: + """Create the declarative object pool and cache robot command states.""" + self.simulation = simulation + self.robot = robot + self.suite = suite + self.track = track + self.control_part = control_part + self._base_full_qpos = robot.get_qpos().clone() + + object_values = track.config.get("objects", []) + if not isinstance(object_values, list): + raise TypeError("atomic-task objects must be a list of mappings.") + for value in object_values: + if not isinstance(value, Mapping): + raise TypeError("Every atomic-task object must be a mapping.") + handle = create_atomic_object(simulation, value) + if handle.object_id in self._objects: + raise ValueError(f"Duplicate atomic object id {handle.object_id!r}.") + self._objects[handle.object_id] = handle + for index, handle in enumerate(self._objects.values()): + handle.park(index) + simulation.update(step=1) + + if any( + isinstance(value, Mapping) and value.get("id") in _GRIPPER_SKILLS + for value in track.config.get("skills", []) + ): + gripper_value = track.config.get("gripper") + if not isinstance(gripper_value, Mapping): + raise ValueError( + "Atomic gripper-action tracks must define a gripper mapping." + ) + end_effector_part = gripper_value.get("control_part") + if not isinstance(end_effector_part, str) or not end_effector_part: + raise ValueError("gripper.control_part must be a non-empty string.") + self.end_effector_part = end_effector_part + limits = robot.get_qpos_limits(name=end_effector_part)[0].to( + device=robot.device, dtype=torch.float32 + ) + dofs = limits.shape[0] + open_qpos = _float_vector( + gripper_value.get("open_qpos"), + name="gripper.open_qpos", + length=dofs, + default=limits[:, 0].detach().cpu().tolist(), + ) + grasp_qpos = _float_vector( + gripper_value.get("grasp_qpos"), + name="gripper.grasp_qpos", + length=dofs, + ) + self._gripper_open = torch.tensor( + open_qpos, dtype=limits.dtype, device=limits.device + ) + self._gripper_grasp = torch.tensor( + grasp_qpos, dtype=limits.dtype, device=limits.device + ) + if bool( + ( + (self._gripper_open < limits[:, 0]) + | (self._gripper_open > limits[:, 1]) + | (self._gripper_grasp < limits[:, 0]) + | (self._gripper_grasp > limits[:, 1]) + ) + .any() + .item() + ): + raise ValueError( + "gripper open/grasp qpos must lie within joint limits." + ) + + def generate_cases( + self, + suite: SuiteCfg, + track: TrackCfg, + robot: Robot, + control_part: str, + batch_size: int, + ) -> list[BenchmarkCase]: + """Generate the algorithm-independent Atomic Task case manifest.""" + if robot is not self.robot or control_part != self.control_part: + raise RuntimeError( + "AtomicTaskScenario must be configured before generation." + ) + skill_values = track.config.get("skills", []) + if not isinstance(skill_values, list) or not skill_values: + raise ValueError("atomic-task skills must be a non-empty list.") + seeds = [int(value) for value in track.config.get("seeds", [11])] + if not seeds: + raise ValueError("atomic-task seeds must not be empty.") + + cases: list[BenchmarkCase] = [] + for skill_index, skill_value in enumerate(skill_values): + if not isinstance(skill_value, Mapping): + raise TypeError("Every atomic-task skill entry must be a mapping.") + skill_id = skill_value.get("id") + if not isinstance(skill_id, str) or not skill_id: + raise ValueError("Every atomic-task skill entry must define an id.") + raw_cases = skill_value.get("cases", []) + if not isinstance(raw_cases, list) or not raw_cases: + raise ValueError(f"Atomic skill {skill_id!r} needs at least one case.") + defaults = { + key: value + for key, value in skill_value.items() + if key not in {"id", "cases"} + } + for case_index, raw_case in enumerate(raw_cases): + if not isinstance(raw_case, Mapping): + raise TypeError("Every atomic skill case must be a mapping.") + config = {**defaults, **dict(raw_case)} + for seed in seeds: + provider = create_atomic_skill_provider(skill_id) + rng_devices = ( + [robot.device] + if torch.device(robot.device).type == "cuda" + else [] + ) + with torch.random.fork_rng(devices=rng_devices): + torch.manual_seed( + _case_generation_seed( + seed, + skill_index=skill_index, + case_index=case_index, + ) + ) + case = provider.generate_case( + self, + suite, + track, + config, + seed=seed, + batch_size=batch_size, + ) + if case.case_id in self._case_providers: + raise ValueError( + f"Duplicate Atomic Task case {case.case_id!r}." + ) + self._case_providers[case.case_id] = provider + cases.append(case) + self.restore_base_robot() + for index, handle in enumerate(self._objects.values()): + handle.park(index) + return cases + + def prepare_planner( + self, adapter: PlannerAdapter, first_case: BenchmarkCase + ) -> None: + """Bind one AtomicActionEngine to the adapter-owned MotionGenerator.""" + del first_case + control_profiles = None + if any(provider.requires_gripper for provider in self._case_providers.values()): + if ( + self.end_effector_part is None + or self._gripper_open is None + or self._gripper_grasp is None + ): + raise RuntimeError("Configured gripper command states are unavailable.") + control_profiles = { + self.end_effector_part: ControlPartCommandProfile.joint_positions( + open=self._gripper_open, + grasp=self._gripper_grasp, + ) + } + motion_generator = adapter.require_motion_generator() + scene_entities = tuple(handle.entity for handle in self._objects.values()) + if scene_entities: + self._engine = create_simulation_atomic_action_engine( + motion_generator=motion_generator, + scene_entities=scene_entities, + control_profiles=control_profiles, + ) + else: + self._engine = AtomicActionEngine( + motion_generator=motion_generator, + control_profiles=control_profiles, + ) + + def close_planner(self, adapter: PlannerAdapter) -> None: + """Drop the engine before its adapter closes the shared generator.""" + del adapter + self._engine = None + + def require_engine(self) -> AtomicActionEngine: + """Return the prepared Atomic Action engine.""" + if self._engine is None: + raise RuntimeError("Atomic Task planner resources were not prepared.") + return self._engine + + def reset_case( + self, + simulation: SimulationManager, + robot: Robot, + case: BenchmarkCase, + control_part: str, + ) -> None: + """Restore full robot and object state before every planner call.""" + del control_part + if case.full_start_qpos is None: + raise ValueError("Atomic Task cases require full_start_qpos.") + for target in (False, True): + robot.set_qpos(case.full_start_qpos, target=target) + robot.clear_dynamics() + active_id = case.object_id + for index, handle in enumerate(self._objects.values()): + if handle.object_id == active_id: + initial_pose = case.case_parameters.get("object_initial_pose") + if initial_pose is None: + handle.reset() + else: + handle.entity.set_local_pose( + torch.tensor( + initial_pose, + dtype=torch.float32, + device=robot.device, + ) + ) + handle.entity.clear_dynamics() + else: + handle.park(index) + simulation.update(step=2) + + def plan_case(self, adapter: PlannerAdapter, case: BenchmarkCase) -> object: + """Compile one Atomic Action with an explicitly pinned motion backend.""" + engine = self.require_engine() + if self.simulation is None: + raise RuntimeError("Atomic Task simulation is not configured.") + provider = self._case_providers[case.case_id] + invocation = provider.build_invocation(self, case, adapter) + policy = invocation.motion_policy + if policy.strategy != "motion_gen": + raise RuntimeError( + "Atomic Task invocations must use the motion_gen strategy." + ) + if engine.motion_generator is not adapter.require_motion_generator(): + raise RuntimeError( + "Atomic Task engine must own the selected adapter's MotionGenerator." + ) + task = provider.initial_task_state(self, case) + context = engine.initial_context( + task=task, + control_dt=float(self.simulation.sim_config.physics_dt), + ) + return engine.compile((invocation,), context=context) + + def plan_contract_error(self, result: object) -> str | None: + """Accept compiled Atomic Action trajectories instead of raw plans.""" + if isinstance(result, CompiledTrajectory): + return None + return f"Expected CompiledTrajectory, got {type(result).__name__}." + + def failure_outcomes( + self, case: BenchmarkCase, failure_code: str + ) -> tuple[CaseOutcome, ...]: + """Mark execution/task stages false after runner-level failures.""" + return tuple( + replace( + outcome, + execution_success=False, + task_success=False, + replan_count=0, + ) + for outcome in super().failure_outcomes(case, failure_code) + ) + + def evaluate_case( + self, + result: object, + case: BenchmarkCase, + robot: Robot, + control_part: str, + suite: SuiteCfg, + *, + planning_time_ms: float, + ) -> ScenarioEvaluation: + """Validate motion, replay physics, and evaluate physical task success.""" + if not isinstance(result, CompiledTrajectory): + raise TypeError(self.plan_contract_error(result)) + arm_joint_ids = list(robot.get_joint_ids(name=control_part)) + arm_positions = result.trajectory.positions[:, :, arm_joint_ids] + motion_outcomes = compute_case_outcomes( + PlanResult( + success=result.plan_success, + positions=arm_positions, + dt=result.trajectory.dt, + ), + case, + robot, + control_part, + validation_samples=suite.protocol.validation_samples, + position_threshold_m=suite.protocol.position_threshold_m, + rotation_threshold_rad=suite.protocol.rotation_threshold_rad, + joint_limit_tolerance_rad=suite.protocol.joint_limit_tolerance_rad, + ) + provider = self._case_providers[case.case_id] + observation = self._execute(result, case, provider) + if observation is None: + execution_success = torch.zeros( + case.batch_size, dtype=torch.bool, device=robot.device + ) + task_success = execution_success.clone() + execution_time_ms = None + tracking = torch.full((case.batch_size,), torch.nan, device=robot.device) + object_lift = None + task_failure_code = "task_goal_miss" + else: + execution_success = observation.execution_success + task_success, task_failure_code = provider.task_result( + self, case, result, observation, motion_outcomes + ) + execution_time_ms = observation.execution_time_ms + tracking = observation.joint_tracking_rmse_rad + object_lift = observation.object_lift_delta_m + + executed_translation_mm: torch.Tensor | None = None + executed_rotation_deg: torch.Tensor | None = None + if observation is not None: + final_target = case.target_waypoints[:, -1] + executed_translation_mm = ( + torch.linalg.vector_norm( + observation.final_tcp_pose[:, :3, 3] - final_target[:, :3, 3], + dim=-1, + ) + * 1000.0 + ) + executed_relative = ( + final_target[:, :3, :3].transpose(-1, -2) + @ observation.final_tcp_pose[:, :3, :3] + ) + executed_trace = torch.diagonal(executed_relative, dim1=-2, dim2=-1).sum( + dim=-1 + ) + executed_rotation_deg = ( + torch.arccos(torch.clamp((executed_trace - 1.0) * 0.5, -1.0, 1.0)) + * 180.0 + / math.pi + ) + + durations = result.trajectory.duration.detach().to("cpu") + outcomes: list[CaseOutcome] = [] + for index, outcome in enumerate(motion_outcomes): + executed = bool(execution_success[index].item()) + task_done = bool(task_success[index].item()) + if outcome.failure_code is not None: + failure_code = outcome.failure_code + elif not executed: + failure_code = "controller_tracking_failure" + elif not task_done: + failure_code = task_failure_code + else: + failure_code = None + tracking_value = float(tracking[index].item()) + outcomes.append( + replace( + outcome, + execution_success=executed, + task_success=task_done, + task_completion_time_s=( + observation.task_completion_time_s + if observation is not None and task_done + else None + ), + joint_tracking_rmse_rad=( + tracking_value if math.isfinite(tracking_value) else None + ), + object_lift_delta_m=( + None + if object_lift is None + else float(object_lift[index].item()) + ), + replan_count=0, + failure_code=failure_code, + executed_final_translation_err_mm=( + None + if executed_translation_mm is None + else float(executed_translation_mm[index].item()) + ), + executed_final_rotation_err_deg=( + None + if executed_rotation_deg is None + else float(executed_rotation_deg[index].item()) + ), + ) + ) + trajectory_duration = ( + float(durations.mean().item()) if durations.numel() else None + ) + return ScenarioEvaluation( + outcomes=tuple(outcomes), + execution_time_ms=execution_time_ms, + end_to_end_time_ms=( + planning_time_ms + execution_time_ms + if execution_time_ms is not None + else None + ), + trajectory_duration_s=trajectory_duration, + trajectory_waypoints=result.trajectory.waypoint_count, + metadata={ + "timing_scope": "atomic_action_compile", + "motion_policy_strategy": "motion_gen", + "physics_validation": "common_joint_target_replay", + "constraint_information": ( + "empty external world; manipulated target excluded from " + "cuRobo collision obstacles" + ), + }, + ) + + def record_replay( + self, + result: object, + case: BenchmarkCase, + evaluation: ScenarioEvaluation | None, + *, + output_dir: Path, + algorithm_id: str, + video: VideoRecordCfg, + ) -> Path | None: + """Reset the case and record a second untimed physics replay.""" + del evaluation + if self.simulation is None or self.robot is None: + return None + self.reset_case(self.simulation, self.robot, case, self.control_part) + compiled = result if isinstance(result, CompiledTrajectory) else None + replayable = compiled is not None and self._is_recordable(compiled) + provider = self._case_providers.get(case.case_id) + video_path = build_video_path( + output_dir, algorithm_id, case.skill_id, case.case_id + ) + + def _replay() -> None: + if replayable and compiled is not None: + self._replay_physics(compiled, case, provider, collect_metrics=False) + else: + self._hold_static() + + return record_with_window(self.simulation, video, video_path, _replay) + + def _execute( + self, + compiled: CompiledTrajectory, + case: BenchmarkCase, + provider: AtomicSkillCaseProvider, + ) -> _ExecutionObservation | None: + """Replay a successful full-robot trajectory under common physics.""" + if self.simulation is None or self.robot is None or self.track is None: + raise RuntimeError("Atomic Task runtime is not configured.") + if not self._is_replayable(compiled): + return None + object_handle = ( + None if case.object_id is None else self.object_handle(case.object_id) + ) + initial_object_position = ( + None + if object_handle is None + else object_handle.entity.get_local_pose(to_matrix=True)[:, :3, 3].clone() + ) + settings = self._physics_settings() + self._synchronize() + started = time.perf_counter() + tracking_parts = self._replay_physics( + compiled, case, provider, collect_metrics=True + ) + self._synchronize() + elapsed_ms = (time.perf_counter() - started) * 1000.0 + if tracking_parts is None: + raise RuntimeError("Timed Atomic Task replay did not return metrics.") + squared_tracking_error, tracking_value_count = tracking_parts + observed = self.robot.get_qpos() + tracking = torch.sqrt(squared_tracking_error / max(tracking_value_count, 1)) + trajectory = compiled.trajectory.positions + simulated_execution_time = ( + trajectory.shape[1] * settings.steps_per_waypoint + + settings.hold_steps * settings.hold_sim_steps + ) * float(self.simulation.sim_config.physics_dt) + final_tcp = self.robot.compute_fk( + self.robot.get_qpos(name=self.control_part), + name=self.control_part, + to_matrix=True, + ) + object_lift = None + final_object_pose = None + if object_handle is not None and initial_object_position is not None: + final_object_pose = object_handle.entity.get_local_pose( + to_matrix=True + ).clone() + final_object_position = final_object_pose[:, :3, 3] + object_lift = final_object_position[:, 2] - initial_object_position[:, 2] + return _ExecutionObservation( + execution_success=torch.isfinite(observed).all(dim=1) + & (tracking <= settings.joint_tracking_tolerance_rad), + final_tcp_pose=final_tcp, + joint_tracking_rmse_rad=tracking, + execution_time_ms=elapsed_ms, + task_completion_time_s=simulated_execution_time, + object_lift_delta_m=object_lift, + final_arm_qpos=self.robot.get_qpos(name=self.control_part).clone(), + final_object_pose=final_object_pose, + ) + + def _physics_settings(self) -> _PhysicsReplaySettings: + """Resolve and validate the common physics-replay step counts.""" + if self.track is None: + raise RuntimeError("Atomic Task runtime is not configured.") + physics = dict(self.track.config.get("physics", {})) + settings = _PhysicsReplaySettings( + steps_per_waypoint=int(physics.get("steps_per_waypoint", 4)), + hold_steps=int(physics.get("hold_steps", 80)), + hold_sim_steps=int(physics.get("hold_sim_steps", 2)), + joint_tracking_tolerance_rad=float( + physics.get("joint_tracking_tolerance_rad", 0.05) + ), + ) + if ( + settings.steps_per_waypoint < 1 + or settings.hold_steps < 0 + or settings.hold_sim_steps < 1 + ): + raise ValueError("Atomic Task physics step counts are invalid.") + return settings + + @staticmethod + def _is_replayable(compiled: CompiledTrajectory) -> bool: + """Return whether a compiled trajectory can be physically replayed.""" + trajectory = compiled.trajectory.positions + return ( + trajectory.shape[1] > 0 + and bool(compiled.plan_success.all().item()) + and bool(torch.isfinite(trajectory).all().item()) + ) + + @staticmethod + def _is_recordable(compiled: CompiledTrajectory) -> bool: + """Allow finite failed planner rollouts in diagnostic videos.""" + trajectory = compiled.trajectory.positions + return trajectory.shape[1] > 0 and bool(torch.isfinite(trajectory).all().item()) + + def _replay_physics( + self, + compiled: CompiledTrajectory, + case: BenchmarkCase, + provider: AtomicSkillCaseProvider | None, + *, + collect_metrics: bool, + ) -> tuple[torch.Tensor, int] | None: + """Replay one compiled trajectory with the evaluation physics contract.""" + if self.simulation is None or self.robot is None: + raise RuntimeError("Atomic Task runtime is not configured.") + trajectory = compiled.trajectory.positions + object_handle = ( + None if case.object_id is None else self.object_handle(case.object_id) + ) + settings = self._physics_settings() + lift_start = None if provider is None else provider.lift_segment_start(compiled) + dynamics_cleared = False + squared_tracking_error: torch.Tensor | None = None + tracking_value_count = 0 + if collect_metrics: + squared_tracking_error = torch.zeros( + case.batch_size, dtype=trajectory.dtype, device=trajectory.device + ) + for waypoint_index in range(trajectory.shape[1]): + positions = trajectory[:, waypoint_index] + self.robot.set_qpos(positions, target=True) + self.simulation.update(step=settings.steps_per_waypoint) + if squared_tracking_error is not None: + observed = self.robot.get_qpos() + squared_tracking_error += ((observed - positions) ** 2).sum(dim=1) + tracking_value_count += observed.shape[1] + if ( + object_handle is not None + and lift_start is not None + and not dynamics_cleared + and waypoint_index + 1 >= lift_start + ): + object_handle.entity.clear_dynamics() + dynamics_cleared = True + final_command = trajectory[:, -1] + for _ in range(settings.hold_steps): + self.robot.set_qpos(final_command, target=True) + self.simulation.update(step=settings.hold_sim_steps) + if squared_tracking_error is not None: + observed = self.robot.get_qpos() + squared_tracking_error += ((observed - final_command) ** 2).sum(dim=1) + tracking_value_count += observed.shape[1] + if collect_metrics: + if squared_tracking_error is None: + raise RuntimeError("Timed replay lost its tracking accumulator.") + return squared_tracking_error, tracking_value_count + return None + + def _hold_static(self) -> None: + """Hold the current scene so a failed-case debug video has frames.""" + if self.simulation is None: + raise RuntimeError("Atomic Task runtime is not configured.") + settings = self._physics_settings() + for _ in range(max(settings.hold_steps, 1)): + self.simulation.update(step=settings.hold_sim_steps) + + def solve_reference_qpos( + self, start_qpos: torch.Tensor, target_waypoints: torch.Tensor + ) -> torch.Tensor: + """Build independent sequential-IK validity evidence for a case.""" + if self.robot is None: + raise RuntimeError("Atomic Task runtime is not configured.") + seed = start_qpos + references: list[torch.Tensor] = [] + for index in range(target_waypoints.shape[1]): + success, seed = self.robot.compute_ik( + pose=target_waypoints[:, index], + joint_seed=seed, + name=self.control_part, + ) + if not bool(torch.as_tensor(success).all().item()): + raise RuntimeError( + f"Independent IK rejected atomic target waypoint {index}." + ) + seed = _canonical_case_qpos(seed) + references.append(seed.clone()) + return torch.stack(references, dim=1) + + def resolve_antipodal_grasp( + self, + handle: AtomicObjectHandle, + object_pose: torch.Tensor, + approach_direction: torch.Tensor, + *, + seed: int, + start_qpos: torch.Tensor, + pre_grasp_distance: float, + lift_height: float, + n_sample: int, + max_candidates: int, + alignment_max_angle_deg: float, + ) -> torch.Tensor: + """Freeze one geometry-aware, independently reachable PGI grasp pose. + + Grasp sampling and sequential IK screening happen once while the case + manifest is built, before any planner adapter is evaluated. Every + planner therefore receives the same explicit grasp pose and planning + latency excludes grasp generation. + """ + if self.robot is None: + raise RuntimeError("Atomic Task runtime is not configured.") + if n_sample < 1 or max_candidates < 1: + raise ValueError( + "Antipodal grasp sample/candidate counts must be positive." + ) + if not 0.0 < alignment_max_angle_deg <= 90.0: + raise ValueError("grasp_alignment_max_angle_deg must be in (0, 90].") + from scripts.tutorials.atomic_action.tutorial_utils import ( + create_antipodal_semantics, + create_parallel_jaw_grasp_pose_generator, + ) + + fork_devices = ( + [] + if self.robot.device.type != "cuda" + else [self.robot.device.index or torch.cuda.current_device()] + ) + with torch.random.fork_rng(devices=fork_devices): + torch.manual_seed(seed) + semantics = create_antipodal_semantics( + handle.entity, + label=handle.object_id, + ) + affordance = semantics.affordance + generator = create_parallel_jaw_grasp_pose_generator( + n_sample=n_sample, + force_refresh=False, + ) + candidates, costs = generator.get_valid_grasp_poses( + mesh_vertices=affordance.mesh_vertices, + mesh_triangles=affordance.mesh_triangles, + obj_poses=object_pose, + approach_direction=approach_direction, + )[0] + if candidates.shape[0] == 0: + raise RuntimeError( + f"No antipodal grasp candidates were found for {handle.object_id!r}." + ) + finite_indices = torch.nonzero(torch.isfinite(costs), as_tuple=False).flatten() + if finite_indices.numel() == 0: + raise RuntimeError( + f"No valid antipodal grasp candidates were found for {handle.object_id!r}." + ) + ranked = finite_indices[torch.argsort(costs[finite_indices])] + ranked = ranked[:max_candidates].detach().to("cpu").tolist() + minimum_alignment = math.cos(math.radians(alignment_max_angle_deg)) + for candidate_index in ranked: + candidate = candidates[candidate_index].to( + device=self.robot.device, dtype=torch.float32 + ) + if ( + float(torch.dot(candidate[:3, 2], approach_direction).item()) + < minimum_alignment + ): + continue + mirrored = candidate.clone() + mirrored[:3, 0] = -mirrored[:3, 0] + mirrored[:3, 1] = -mirrored[:3, 1] + for variant in (candidate, mirrored): + grasp = variant.unsqueeze(0) + pre_grasp = grasp.clone() + pre_grasp[:, :3, 3] -= approach_direction * pre_grasp_distance + lift = grasp.clone() + lift[:, 2, 3] += lift_height + seed = start_qpos + feasible = True + for pose in (pre_grasp, grasp, lift): + success, seed = self.robot.compute_ik( + pose=pose, + joint_seed=seed, + name=self.control_part, + ) + if not bool(torch.as_tensor(success).all().item()): + feasible = False + break + seed = _canonical_case_qpos(seed) + if feasible: + return grasp + raise RuntimeError( + f"No independently reachable antipodal grasp remained for " + f"{handle.object_id!r} after screening {len(ranked)} candidates." + ) + + def restore_base_robot(self) -> None: + """Restore the robot state captured before scenario case generation.""" + if self.robot is None or self._base_full_qpos is None: + raise RuntimeError("Atomic Task runtime is not configured.") + for target in (False, True): + self.robot.set_qpos(self._base_full_qpos, target=target) + self.robot.clear_dynamics() + + def randomize_robot_start( + self, + config: Mapping[str, object], + *, + seed: int, + stream: int, + ) -> torch.Tensor: + """Apply a deterministic, bounded perturbation to the arm start state.""" + if self.robot is None: + raise RuntimeError("Atomic Task runtime is not configured.") + start = self.robot.get_qpos(name=self.control_part).clone() + amplitude = _float_vector( + config.get("start_qpos_jitter_rad"), + name="start_qpos_jitter_rad", + length=start.shape[-1], + default=(0.0,) * start.shape[-1], + ) + if any(value < 0.0 for value in amplitude): + raise ValueError("start_qpos_jitter_rad values must be non-negative.") + randomized = start + _seeded_jitter( + amplitude, + seed=seed, + stream=stream, + dtype=start.dtype, + device=start.device, + ) + limits = self.robot.get_qpos_limits(name=self.control_part)[0] + margin = float(config.get("start_joint_limit_margin_rad", 0.05)) + if not math.isfinite(margin) or margin < 0.0: + raise ValueError("start_joint_limit_margin_rad must be non-negative.") + lower = limits[:, 0] + margin + upper = limits[:, 1] - margin + if bool(((randomized < lower) | (randomized > upper)).any().item()): + raise RuntimeError( + "Randomized arm start violates the configured joint-limit margin." + ) + self.set_robot_start( + randomized, + open_gripper=False, + grasp_gripper=False, + ) + return randomized + + def randomize_object_pose( + self, + handle: AtomicObjectHandle, + config: Mapping[str, object], + *, + seed: int, + stream: int, + ) -> torch.Tensor: + """Install one deterministic position/yaw perturbation for an object.""" + pose = handle.initial_pose.clone() + position_amplitude = _float_vector( + config.get("object_position_jitter_m"), + name="object_position_jitter_m", + length=3, + default=(0.0, 0.0, 0.0), + ) + if any(value < 0.0 for value in position_amplitude): + raise ValueError("object_position_jitter_m values must be non-negative.") + pose[:, :3, 3] += _seeded_jitter( + position_amplitude, + seed=seed, + stream=stream, + dtype=pose.dtype, + device=pose.device, + ) + + yaw_amplitude = float(config.get("object_yaw_jitter_rad", 0.0)) + if not math.isfinite(yaw_amplitude) or yaw_amplitude < 0.0: + raise ValueError("object_yaw_jitter_rad must be finite and non-negative.") + yaw = _seeded_jitter( + (yaw_amplitude,), + seed=seed, + stream=stream + 1, + dtype=pose.dtype, + device=pose.device, + )[0] + cosine = torch.cos(yaw) + sine = torch.sin(yaw) + yaw_rotation = torch.eye(3, dtype=pose.dtype, device=pose.device) + yaw_rotation[0, 0] = cosine + yaw_rotation[0, 1] = -sine + yaw_rotation[1, 0] = sine + yaw_rotation[1, 1] = cosine + pose[:, :3, :3] = yaw_rotation.unsqueeze(0) @ pose[:, :3, :3] + + handle.entity.set_local_pose(pose) + handle.entity.clear_dynamics() + if self.simulation is not None: + self.simulation.update(step=2) + handle.entity.set_local_pose(pose) + handle.entity.clear_dynamics() + return pose + + def set_robot_start( + self, + manipulator_qpos: torch.Tensor, + *, + open_gripper: bool, + grasp_gripper: bool = False, + ) -> None: + """Install a manipulator start and optional gripper command.""" + if self.robot is None: + raise RuntimeError("Atomic Task runtime is not configured.") + if open_gripper and grasp_gripper: + raise ValueError("The gripper cannot start both open and grasping.") + for target in (False, True): + self.robot.set_qpos(manipulator_qpos, name=self.control_part, target=target) + if open_gripper or grasp_gripper: + gripper_qpos = ( + self._gripper_open if open_gripper else self._gripper_grasp + ) + if self.end_effector_part is None or gripper_qpos is None: + raise RuntimeError("Requested gripper state is unavailable.") + command = gripper_qpos.unsqueeze(0).expand( + manipulator_qpos.shape[0], -1 + ) + self.robot.set_qpos(command, name=self.end_effector_part, target=target) + self.robot.clear_dynamics() + + def activate_object(self, object_id: str) -> AtomicObjectHandle: + """Reset one case object and park every other configured object.""" + handle = self.object_handle(object_id) + for index, candidate in enumerate(self._objects.values()): + if candidate is handle: + candidate.reset() + else: + candidate.park(index) + if self.simulation is not None: + self.simulation.update(step=2) + return handle + + def object_handle(self, object_id: str | None) -> AtomicObjectHandle: + """Resolve a configured object identifier with an actionable error.""" + if object_id is None: + raise ValueError("This Atomic Task case has no object id.") + try: + return self._objects[object_id] + except KeyError as exc: + raise ValueError( + f"Unknown atomic object {object_id!r}; configured objects: " + f"{sorted(self._objects)}." + ) from exc + + def close_runtime(self) -> None: + """Release Python references before SimulationManager teardown.""" + self._engine = None + self._case_providers.clear() + self._objects.clear() + self._base_full_qpos = None + self.end_effector_part = None + self._gripper_open = None + self._gripper_grasp = None + self.simulation = None + self.robot = None + self.suite = None + self.track = None + + @staticmethod + def _synchronize() -> None: + """Synchronize CUDA around physical wall-time measurement.""" + if torch.cuda.is_available(): + torch.cuda.synchronize() + + +register_atomic_skill_provider("move_end_effector", _MoveEndEffectorCases) +register_atomic_skill_provider("move_held_object", _MoveHeldObjectCases) +register_atomic_skill_provider("move_joints", _MoveJointsCases) +register_atomic_skill_provider("pick_up", _PickUpCases) +register_atomic_skill_provider("place", _PlaceCases) +register_atomic_skill_provider("press", _PressCases) +register_scenario_provider("atomic_task", AtomicTaskScenario) diff --git a/scripts/benchmark/motion_generation/scenarios/base.py b/scripts/benchmark/motion_generation/scenarios/base.py index 11435cb71..4be5b3c72 100644 --- a/scripts/benchmark/motion_generation/scenarios/base.py +++ b/scripts/benchmark/motion_generation/scenarios/base.py @@ -19,19 +19,41 @@ from __future__ import annotations from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from pathlib import Path from typing import TYPE_CHECKING +from embodichain.lab.sim.planners.utils import PlanResult + +from ..metrics.trajectory import compute_case_outcomes, make_failure_outcomes +from ..models import CaseOutcome + if TYPE_CHECKING: + from embodichain.lab.sim import SimulationManager from embodichain.lab.sim.objects import Robot from ..config import SuiteCfg, TrackCfg from ..models import BenchmarkCase + from ..planners.base import PlannerAdapter + from ..video import VideoRecordCfg + +__all__ = ["ScenarioEvaluation", "ScenarioProvider"] + -__all__ = ["ScenarioProvider"] +@dataclass(frozen=True) +class ScenarioEvaluation: + """Outcomes and execution metrics produced outside planner timing.""" + + outcomes: tuple[CaseOutcome, ...] + execution_time_ms: float | None = None + end_to_end_time_ms: float | None = None + trajectory_duration_s: float | None = None + trajectory_waypoints: int | None = None + metadata: dict[str, object] = field(default_factory=dict) class ScenarioProvider(ABC): - """Generate fixed cases for one registered scenario kind.""" + """Generate, plan, execute, and evaluate one registered scenario kind.""" required_capabilities: frozenset[str] = frozenset() @@ -49,3 +71,108 @@ def generate_cases( batch_size: int, ) -> list["BenchmarkCase"]: """Build the frozen case manifest for one batch size.""" + + def configure_runtime( + self, + simulation: "SimulationManager", + robot: "Robot", + suite: "SuiteCfg", + track: "TrackCfg", + control_part: str, + ) -> None: + """Create scenario-owned runtime entities before case generation.""" + + def close_runtime(self) -> None: + """Release references to scenario-owned simulation entities.""" + + def prepare_planner( + self, adapter: "PlannerAdapter", first_case: "BenchmarkCase" + ) -> None: + """Bind scenario resources to a built planner outside trial timing.""" + + def close_planner(self, adapter: "PlannerAdapter") -> None: + """Release scenario resources that retain a planner adapter.""" + + def reset_case( + self, + simulation: "SimulationManager", + robot: "Robot", + case: "BenchmarkCase", + control_part: str, + ) -> None: + """Restore the frozen robot start state before a planning call.""" + if case.full_start_qpos is not None: + for target in (False, True): + robot.set_qpos(case.full_start_qpos, target=target) + else: + for target in (False, True): + robot.set_qpos(case.start_qpos, name=control_part, target=target) + robot.clear_dynamics() + simulation.update(step=1) + + def plan_case(self, adapter: "PlannerAdapter", case: "BenchmarkCase") -> object: + """Plan one case through the selected adapter.""" + return adapter.plan(case) + + def plan_contract_error(self, result: object) -> str | None: + """Return a diagnostic when a planner artifact violates this scenario.""" + if isinstance(result, PlanResult): + return None + return f"Expected PlanResult, got {type(result).__name__}." + + def failure_outcomes( + self, case: "BenchmarkCase", failure_code: str + ) -> tuple[CaseOutcome, ...]: + """Create scenario-appropriate outcomes after a runner-level failure.""" + return make_failure_outcomes(case.batch_size, failure_code) + + def evaluate_case( + self, + result: object, + case: "BenchmarkCase", + robot: "Robot", + control_part: str, + suite: "SuiteCfg", + *, + planning_time_ms: float, + ) -> ScenarioEvaluation: + """Externally validate a planner-only trajectory.""" + if not isinstance(result, PlanResult): + raise TypeError(self.plan_contract_error(result)) + outcomes = compute_case_outcomes( + result, + case, + robot, + control_part, + validation_samples=suite.protocol.validation_samples, + position_threshold_m=suite.protocol.position_threshold_m, + rotation_threshold_rad=suite.protocol.rotation_threshold_rad, + joint_limit_tolerance_rad=suite.protocol.joint_limit_tolerance_rad, + ) + return ScenarioEvaluation(outcomes=outcomes) + + def record_replay( + self, + result: object, + case: "BenchmarkCase", + evaluation: ScenarioEvaluation | None, + *, + output_dir: Path, + algorithm_id: str, + video: "VideoRecordCfg", + ) -> Path | None: + """Optionally record a second, untimed replay. Default is a no-op. + + Args: + result: Planner or compiled-action artifact from the measured trial. + case: Frozen case identity used for the output filename. + evaluation: Timed evaluation, or ``None`` after a runner-level failure. + output_dir: Directory that should receive the mp4. + algorithm_id: Planner id used in the filename. + video: Recording policy and encoder settings. + + Returns: + Path to a saved video, or ``None`` when this scenario does not record. + """ + del result, case, evaluation, output_dir, algorithm_id, video + return None diff --git a/scripts/benchmark/motion_generation/scenarios/free_space.py b/scripts/benchmark/motion_generation/scenarios/free_space.py index ce2b7fff4..f299d7fe6 100644 --- a/scripts/benchmark/motion_generation/scenarios/free_space.py +++ b/scripts/benchmark/motion_generation/scenarios/free_space.py @@ -199,6 +199,7 @@ def _build_case( start_qpos=start_qpos, target_waypoints=target_waypoints, reference_qpos=reference_qpos, + robot_id=suite.robot.id, ) diff --git a/scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo.yaml b/scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo.yaml new file mode 100644 index 000000000..0d3a53396 --- /dev/null +++ b/scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo.yaml @@ -0,0 +1,173 @@ +schema_version: 1 +name: atomic_skill_franka_pgi_curobo +suite_version: atomic_franka_pgi_curobo_smoke_v1 +profile: smoke + +robot: + id: franka_pgi + provider: franka_pgi + config: {} + +planners: + - id: nmg + adapter: nmg_onnx + role: candidate + enabled: false + config: + onnx_model_path: null + num_waypoints: 3 + max_steps: 240 + action_scale: 0.2 + use_relative_obs: true + intermediate_orientation: true + pos_eps: 0.01 + rot_eps: 0.1 + joint_eps: 0.02 + policy_frame_from_world: + - [-1.0, 0.0, 0.0, 0.0] + - [0.0, -1.0, 0.0, 0.0] + - [0.0, 0.0, 1.0, 0.0] + - [0.0, 0.0, 0.0, 1.0] + runtime_tcp_from_policy_tcp: + - [0.70710678, 0.70710678, 0.0, 0.0] + - [-0.70710678, 0.70710678, 0.0, 0.0] + - [0.0, 0.0, 1.0, -0.0466] + - [0.0, 0.0, 0.0, 1.0] + - id: curobo + adapter: curobo + role: primary_baseline + enabled: true + config: + max_attempts: 5 + max_planning_time: null + interpolation_dt: 0.025 + collision_activation_distance: 0.01 + use_cuda_graph: true + cuda_graph_fallback: true + warmup_iterations: 1 + preserve_plan_samples: false + world: + obstacle_representation: sphere + multi_env: false + auto_gen: + fit_type: voxel + sphere_density: 0.1 + collision_sphere_buffer: 0.0 + +protocol: + warmup_trials: 1 + measured_trials: 1 + sample_interval: 80 + validation_samples: 128 + position_threshold_m: 0.01 + rotation_threshold_rad: 0.1 + joint_limit_tolerance_rad: 0.00001 + +tracks: + - id: atomic-task + scenario: atomic_task + enabled: true + config: + batch_sizes: [1] + seeds: [11] + physics: + steps_per_waypoint: 4 + hold_steps: 80 + hold_sim_steps: 2 + joint_tracking_tolerance_rad: 0.05 + gripper: + control_part: hand + open_qpos: [0.0] + grasp_qpos: [0.024] + objects: + - id: cube_50mm + kind: cube + size: [0.05, 0.05, 0.05] + position: [-0.42, -0.08, 0.05] + mass: 0.05 + dynamic_friction: 0.97 + static_friction: 0.99 + contact_offset: 0.003 + rest_offset: 0.001 + settle_steps: 10 + skills: + - id: move_end_effector + sample_count: 80 + cases: + - name: relative_two_waypoint + task_difficulty: simple + target_offsets_m: + - [-0.08, -0.08, 0.08] + - [0.04, 0.12, 0.04] + difficulty_factors: + waypoint_count: 2 + maximum_translation_m: 0.13 + obstacles: 0 + - id: move_joints + sample_count: 80 + joint_limit_margin_rad: 0.05 + joint_threshold_rad: 0.02 + cases: + - name: relative_multi_joint + task_difficulty: simple + target_offsets_rad: + - [0.12, -0.08, 0.10, 0.12, -0.08, 0.10, -0.10] + difficulty_factors: + waypoint_count: 1 + active_joints: 7 + - id: pick_up + sample_count: 120 + grasp_source: antipodal + grasp_sample_count: 10000 + grasp_max_candidates: 128 + grasp_alignment_max_angle_deg: 10.0 + hand_interp_steps: 12 + pre_grasp_distance_m: 0.15 + lift_height_m: 0.16 + minimum_object_lift_m: 0.04 + approach_direction: [0.0, 0.0, -1.0] + pre_pick_height_m: 0.36 + cases: + - name: cube_top_center + object: cube_50mm + task_difficulty: simple + grasp_offset_m: [0.0, 0.0, 0.0] + difficulty_factors: + approach: top + pre_grasp_clearance_m: 0.15 + required_lift_m: 0.04 + - id: move_held_object + sample_count: 80 + held_object_offset_m: [0.0, 0.0, 0.18] + target_object_offset_m: [0.08, 0.08, 0.04] + object_position_threshold_m: 0.04 + cases: + - name: cube_transport + object: cube_50mm + task_difficulty: simple + difficulty_factors: + translation_m: 0.12 + object_preheld: true + - id: place + sample_count: 120 + held_object_offset_m: [0.0, 0.0, 0.18] + target_object_offset_m: [0.10, 0.10, 0.0] + retract_height_m: 0.10 + hand_interp_steps: 12 + object_position_threshold_m: 0.05 + cases: + - name: cube_table_place + object: cube_50mm + task_difficulty: simple + difficulty_factors: + translation_m: 0.14 + object_preheld: true + - id: press + sample_count: 80 + hand_interp_steps: 8 + cases: + - name: vertical_press + task_difficulty: simple + target_offset_m: [0.0, 0.0, -0.08] + difficulty_factors: + press_depth_m: 0.08 diff --git a/scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo_randomized.yaml b/scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo_randomized.yaml new file mode 100644 index 000000000..f440671d7 --- /dev/null +++ b/scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo_randomized.yaml @@ -0,0 +1,188 @@ +schema_version: 1 +name: atomic_skill_franka_pgi_curobo_randomized +suite_version: atomic_franka_pgi_curobo_randomized_v1 +profile: coverage + +robot: + id: franka_pgi + provider: franka_pgi + config: {} + +planners: + - id: nmg + adapter: nmg_onnx + role: candidate + enabled: false + config: + onnx_model_path: null + num_waypoints: 3 + max_steps: 240 + action_scale: 0.2 + use_relative_obs: true + intermediate_orientation: true + pos_eps: 0.01 + rot_eps: 0.1 + joint_eps: 0.02 + policy_frame_from_world: + - [-1.0, 0.0, 0.0, 0.0] + - [0.0, -1.0, 0.0, 0.0] + - [0.0, 0.0, 1.0, 0.0] + - [0.0, 0.0, 0.0, 1.0] + runtime_tcp_from_policy_tcp: + - [0.70710678, 0.70710678, 0.0, 0.0] + - [-0.70710678, 0.70710678, 0.0, 0.0] + - [0.0, 0.0, 1.0, -0.0466] + - [0.0, 0.0, 0.0, 1.0] + - id: curobo + adapter: curobo + role: primary_baseline + enabled: true + config: + max_attempts: 5 + max_planning_time: null + interpolation_dt: 0.025 + collision_activation_distance: 0.01 + use_cuda_graph: true + cuda_graph_fallback: true + warmup_iterations: 1 + preserve_plan_samples: false + world: + obstacle_representation: sphere + multi_env: false + auto_gen: + fit_type: voxel + sphere_density: 0.1 + collision_sphere_buffer: 0.0 + +protocol: + warmup_trials: 0 + measured_trials: 1 + sample_interval: 80 + validation_samples: 128 + position_threshold_m: 0.01 + rotation_threshold_rad: 0.1 + joint_limit_tolerance_rad: 0.00001 + +tracks: + - id: atomic-task-randomized + scenario: atomic_task + enabled: true + config: + batch_sizes: [1] + seeds: [101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116] + physics: + steps_per_waypoint: 4 + hold_steps: 80 + hold_sim_steps: 2 + joint_tracking_tolerance_rad: 0.05 + gripper: + control_part: hand + open_qpos: [0.0] + grasp_qpos: [0.024] + objects: + - id: cube_50mm + kind: cube + size: [0.05, 0.05, 0.05] + position: [-0.42, -0.08, 0.05] + mass: 0.05 + dynamic_friction: 0.97 + static_friction: 0.99 + contact_offset: 0.003 + rest_offset: 0.001 + settle_steps: 10 + skills: + - id: move_end_effector + sample_count: 80 + start_qpos_jitter_rad: [0.15, 0.15, 0.15, 0.12, 0.15, 0.12, 0.15] + target_offset_jitter_m: [0.03, 0.03, 0.02] + cases: + - name: randomized_two_waypoint + task_difficulty: medium + target_offsets_m: + - [-0.08, -0.08, 0.08] + - [0.04, 0.12, 0.04] + difficulty_factors: + waypoint_count: 2 + randomized_start: true + randomized_targets: true + - id: move_joints + sample_count: 80 + joint_limit_margin_rad: 0.05 + joint_threshold_rad: 0.02 + start_qpos_jitter_rad: [0.15, 0.15, 0.15, 0.12, 0.15, 0.12, 0.15] + target_offset_jitter_rad: [0.05, 0.05, 0.05, 0.05, 0.05, 0.05, 0.05] + cases: + - name: randomized_multi_joint + task_difficulty: medium + target_offsets_rad: + - [0.12, -0.08, 0.10, 0.12, -0.08, 0.10, -0.10] + difficulty_factors: + waypoint_count: 1 + active_joints: 7 + randomized_start: true + randomized_target: true + - id: pick_up + sample_count: 120 + grasp_source: fixed + hand_interp_steps: 12 + pre_grasp_distance_m: 0.15 + lift_height_m: 0.16 + minimum_object_lift_m: 0.04 + approach_direction: [0.0, 0.0, -1.0] + pre_pick_height_m: 0.36 + object_position_jitter_m: [0.05, 0.05, 0.0] + object_yaw_jitter_rad: 0.40 + cases: + - name: randomized_cube_pick + object: cube_50mm + task_difficulty: medium + grasp_offset_m: [0.0, 0.0, 0.0] + difficulty_factors: + approach: top + randomized_object_pose: true + - id: move_held_object + sample_count: 80 + held_object_offset_m: [0.0, 0.0, 0.18] + held_object_offset_jitter_m: [0.02, 0.02, 0.015] + target_object_offset_m: [0.08, 0.08, 0.04] + target_object_offset_jitter_m: [0.04, 0.04, 0.02] + object_position_jitter_m: [0.04, 0.04, 0.0] + object_yaw_jitter_rad: 0.40 + object_position_threshold_m: 0.04 + cases: + - name: randomized_cube_transport + object: cube_50mm + task_difficulty: medium + difficulty_factors: + object_preheld: true + randomized_start_and_target: true + - id: place + sample_count: 120 + held_object_offset_m: [0.0, 0.0, 0.18] + held_object_offset_jitter_m: [0.02, 0.02, 0.015] + target_object_offset_m: [0.10, 0.10, 0.0] + target_object_offset_jitter_m: [0.04, 0.04, 0.0] + object_position_jitter_m: [0.04, 0.04, 0.0] + object_yaw_jitter_rad: 0.40 + retract_height_m: 0.10 + hand_interp_steps: 12 + object_position_threshold_m: 0.05 + cases: + - name: randomized_cube_place + object: cube_50mm + task_difficulty: medium + difficulty_factors: + object_preheld: true + randomized_start_and_target: true + - id: press + sample_count: 80 + hand_interp_steps: 8 + start_qpos_jitter_rad: [0.15, 0.15, 0.15, 0.12, 0.15, 0.12, 0.15] + target_offset_jitter_m: [0.03, 0.03, 0.015] + cases: + - name: randomized_vertical_press + task_difficulty: medium + target_offset_m: [0.0, 0.0, -0.08] + difficulty_factors: + randomized_start: true + randomized_target: true diff --git a/scripts/benchmark/motion_generation/suites/coverage.yaml b/scripts/benchmark/motion_generation/suites/coverage.yaml index 6cc331254..fb8b2cf0d 100644 --- a/scripts/benchmark/motion_generation/suites/coverage.yaml +++ b/scripts/benchmark/motion_generation/suites/coverage.yaml @@ -38,13 +38,18 @@ planners: velocity: 0.2 acceleration: 0.5 - id: nmg - adapter: neural_stub + adapter: nmg_onnx role: candidate enabled: false config: - model_revision: not-ready - pos_eps: 0.05 - rot_eps: 0.3 + onnx_model_path: null + num_waypoints: 8 + max_steps: 240 + action_scale: 0.2 + use_relative_obs: true + intermediate_orientation: true + pos_eps: 0.01 + rot_eps: 0.1 protocol: warmup_trials: 3 diff --git a/scripts/benchmark/motion_generation/suites/smoke.yaml b/scripts/benchmark/motion_generation/suites/smoke.yaml index 459fe3f49..60a98c0e2 100644 --- a/scripts/benchmark/motion_generation/suites/smoke.yaml +++ b/scripts/benchmark/motion_generation/suites/smoke.yaml @@ -38,13 +38,18 @@ planners: velocity: 0.2 acceleration: 0.5 - id: nmg - adapter: neural_stub + adapter: nmg_onnx role: candidate enabled: false config: - model_revision: not-ready - pos_eps: 0.05 - rot_eps: 0.3 + onnx_model_path: null + num_waypoints: 8 + max_steps: 240 + action_scale: 0.2 + use_relative_obs: true + intermediate_orientation: true + pos_eps: 0.01 + rot_eps: 0.1 protocol: warmup_trials: 1 diff --git a/scripts/benchmark/motion_generation/video.py b/scripts/benchmark/motion_generation/video.py new file mode 100644 index 000000000..30c140b12 --- /dev/null +++ b/scripts/benchmark/motion_generation/video.py @@ -0,0 +1,198 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Optional Atomic Task replay recording helpers. + +Recording is a second, untimed physics pass. Failures never change trial +success and never add a fourth Markdown table. +""" + +from __future__ import annotations + +import argparse +import re +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from embodichain.lab.sim import SimulationManager + +__all__ = [ + "DEFAULT_VIDEO_LOOK_AT", + "VideoRecordCfg", + "build_video_path", + "record_with_window", + "should_record_case", + "summarize_video_recording", + "video_cfg_from_args", +] + +DEFAULT_VIDEO_FPS = 20 +DEFAULT_VIDEO_MAX_MEMORY_MB = 2048 +DEFAULT_VIDEO_WIDTH = 640 +DEFAULT_VIDEO_HEIGHT = 480 +DEFAULT_VIDEO_CASE_LIMIT = 0 +DEFAULT_VIDEO_LOOK_AT = ( + (-1.25, -1.15, 0.95), + (-0.25, -0.02, 0.25), + (0.0, 0.0, 1.0), +) + +LookAt = tuple[Sequence[float], Sequence[float], Sequence[float]] + + +@dataclass(frozen=True) +class VideoRecordCfg: + """CLI-resolved recording policy for Atomic Task measured replays.""" + + enabled: bool = False + record_failed: bool = False + case_limit: int = DEFAULT_VIDEO_CASE_LIMIT + fps: int = DEFAULT_VIDEO_FPS + width: int = DEFAULT_VIDEO_WIDTH + height: int = DEFAULT_VIDEO_HEIGHT + max_memory_mb: int = DEFAULT_VIDEO_MAX_MEMORY_MB + look_at: LookAt = DEFAULT_VIDEO_LOOK_AT + output_dir: Path | None = None + + def __post_init__(self) -> None: + if self.case_limit < 0: + raise ValueError("video case_limit must be non-negative.") + if self.fps <= 0: + raise ValueError("video fps must be positive.") + if self.width <= 0 or self.height <= 0: + raise ValueError("video width and height must be positive.") + if self.max_memory_mb <= 0: + raise ValueError("video max_memory_mb must be positive.") + + +def should_record_case(cfg: VideoRecordCfg, recorded_count: int, success: bool) -> bool: + """Return whether one measured Atomic Task case should emit a video.""" + if not cfg.enabled: + return False + if not success and not cfg.record_failed: + return False + return cfg.case_limit == 0 or recorded_count < cfg.case_limit + + +def build_video_path( + output_dir: Path, + algorithm_id: str, + skill_id: str, + case_id: str, +) -> Path: + """Build ``{algorithm}_{skill}_{case}.mp4`` under the run videos directory.""" + + def _sanitize(value: str) -> str: + return re.sub(r"[^A-Za-z0-9_.-]+", "_", value).strip("_") or "unknown" + + output_dir.mkdir(parents=True, exist_ok=True) + filename = ( + f"{_sanitize(algorithm_id)}_{_sanitize(skill_id)}_{_sanitize(case_id)}.mp4" + ) + return output_dir / filename + + +def record_with_window( + sim: "SimulationManager", + cfg: VideoRecordCfg, + video_path: Path, + replay_fn: Callable[[], None], +) -> Path | None: + """Record one headless replay. Failures print a warning and return ``None``.""" + try: + original_width = sim.sim_config.width + original_height = sim.sim_config.height + recording_started = False + try: + sim.sim_config.width = cfg.width + sim.sim_config.height = cfg.height + recording_started = sim.start_window_record( + save_path=str(video_path), + fps=cfg.fps, + max_memory=cfg.max_memory_mb, + look_at=cfg.look_at, + use_sim_time=True, + ) + finally: + sim.sim_config.width = original_width + sim.sim_config.height = original_height + if not recording_started: + return None + + stop_success = False + try: + replay_fn() + finally: + if sim.is_window_recording(): + stop_success = sim.stop_window_record() + sim.wait_window_record_saves() + return video_path if stop_success else None + except Exception as exc: # noqa: BLE001 - recording must not change success + try: + if sim.is_window_recording(): + sim.stop_window_record() + sim.wait_window_record_saves() + except Exception: + pass + print( + "Warning: failed to record Atomic Task replay video " + f"{video_path}: {type(exc).__name__}: {exc}" + ) + return None + + +def summarize_video_recording( + cfg: VideoRecordCfg, video_paths: Sequence[str] +) -> list[str]: + """Return report notes describing recording coverage without extra tables.""" + if not cfg.enabled: + return ["Video policy: disabled."] + if cfg.record_failed: + policy = ( + "records Atomic Task measured success replays and failed-case " + "static scenes when capture is available." + ) + else: + policy = ( + "records Atomic Task measured success replays only; failed " + "cases are reported in the tables but do not emit videos." + ) + rendered = ", ".join(video_paths) if video_paths else "none" + return [ + f"Video policy: {policy}", + f"videos={len(video_paths)}", + f"Replay videos: {rendered}", + ] + + +def video_cfg_from_args(args: argparse.Namespace) -> VideoRecordCfg: + """Build recording config from the motion-generation CLI namespace.""" + output_dir = getattr(args, "video_dir", None) + return VideoRecordCfg( + enabled=bool(getattr(args, "record_video", False)), + record_failed=bool(getattr(args, "record_failed_video", False)), + case_limit=int(getattr(args, "video_case_limit", DEFAULT_VIDEO_CASE_LIMIT)), + fps=int(getattr(args, "video_fps", DEFAULT_VIDEO_FPS)), + width=int(getattr(args, "video_width", DEFAULT_VIDEO_WIDTH)), + height=int(getattr(args, "video_height", DEFAULT_VIDEO_HEIGHT)), + max_memory_mb=int( + getattr(args, "video_max_memory", DEFAULT_VIDEO_MAX_MEMORY_MB) + ), + output_dir=None if output_dir in (None, "") else Path(output_dir), + ) diff --git a/tests/benchmark/motion_generation/test_atomic_task_benchmark.py b/tests/benchmark/motion_generation/test_atomic_task_benchmark.py new file mode 100644 index 000000000..c43309b07 --- /dev/null +++ b/tests/benchmark/motion_generation/test_atomic_task_benchmark.py @@ -0,0 +1,511 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Pure-logic tests for the planner-oriented Atomic Task benchmark track.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from unittest.mock import Mock + +import pytest +import torch + +from scripts.benchmark.motion_generation.aggregation import aggregate_results +from scripts.benchmark.motion_generation.artifacts import write_case_manifest +from scripts.benchmark.motion_generation.config import load_suite +from scripts.benchmark.motion_generation.models import ( + AlgorithmRole, + BenchmarkCase, + CaseOutcome, + PlannerMetadata, + TrialPhase, + TrialRecord, +) +from scripts.benchmark.motion_generation.registry import create_robot_provider +from scripts.benchmark.motion_generation.runner import BenchmarkRunner +from scripts.benchmark.motion_generation import robots as _robots # noqa: F401 +from scripts.benchmark.motion_generation.scenarios.atomic_objects import ( + atomic_object_kind_names, + create_atomic_object, +) +from scripts.benchmark.motion_generation.scenarios.atomic_task import ( + AtomicTaskScenario, + _canonical_case_qpos, + _case_generation_seed, + _randomization_parameters, + _seeded_jitter, + atomic_skill_provider_names, + create_atomic_skill_provider, +) +from scripts.benchmark.motion_generation.scenarios.base import ScenarioProvider +from scripts.benchmark.motion_generation.scenarios.free_space import FreeSpaceScenario +from scripts.benchmark.motion_generation.video import ( + VideoRecordCfg, + build_video_path, + record_with_window, + should_record_case, + summarize_video_recording, + video_cfg_from_args, +) + + +def _atomic_case() -> BenchmarkCase: + target = torch.eye(4).reshape(1, 1, 4, 4) + return BenchmarkCase( + suite_version="atomic_test_v1", + track="atomic-task", + scenario_id="move_end_effector", + case_id="atomic-task:move_end_effector:simple:s11", + seed=11, + batch_size=1, + num_waypoints=1, + path_shape="robot_relative_waypoints", + start_state_bin="pre_action", + start_qpos=torch.zeros(1, 7), + target_waypoints=target, + reference_qpos=torch.zeros(1, 1, 7), + robot_id="franka_pgi", + skill_id="move_end_effector", + task_difficulty="simple", + primary_success="task_success", + full_start_qpos=torch.zeros(1, 9), + case_parameters={"sample_count": 80, "target_offsets_m": [[0.1, 0.0, 0.0]]}, + ) + + +def _atomic_outcome() -> CaseOutcome: + return CaseOutcome( + env_index=0, + planning_success=True, + finite=True, + ordered_waypoints_reached=True, + motion_valid=True, + completed_waypoint_ratio=1.0, + final_translation_err_mm=1.0, + final_rotation_err_deg=1.0, + waypoint_translation_err_mm_mean=1.0, + waypoint_translation_err_mm_p95=1.0, + waypoint_translation_err_mm_max=1.0, + waypoint_rotation_err_deg_mean=1.0, + waypoint_rotation_err_deg_p95=1.0, + waypoint_rotation_err_deg_max=1.0, + joint_limit_violation=False, + max_normalized_joint_violation=0.0, + joint_path_length_rad=0.2, + cartesian_path_length_m=0.1, + path_efficiency=1.0, + execution_success=True, + task_success=True, + task_completion_time_s=1.5, + joint_tracking_rmse_rad=0.002, + replan_count=0, + ) + + +def test_atomic_suite_is_franka_pgi_and_curobo_only(): + suite = load_suite("atomic_franka_pgi_curobo") + + assert suite.robot.id == "franka_pgi" + assert suite.robot.provider == "franka_pgi" + assert [spec.id for spec in suite.planners if spec.enabled] == ["curobo"] + assert [(track.id, track.scenario) for track in suite.enabled_tracks()] == [ + ("atomic-task", "atomic_task") + ] + skills = suite.enabled_tracks()[0].config["skills"] + assert [item["id"] for item in skills] == [ + "move_end_effector", + "move_joints", + "pick_up", + "move_held_object", + "place", + "press", + ] + gripper = suite.enabled_tracks()[0].config["gripper"] + assert gripper == { + "control_part": "hand", + "open_qpos": [0.0], + "grasp_qpos": [0.024], + } + + +def test_randomized_atomic_suite_covers_six_skills_with_fixed_seed_sweep(): + suite = load_suite("atomic_franka_pgi_curobo_randomized") + track = suite.enabled_tracks()[0] + + assert suite.suite_version == "atomic_franka_pgi_curobo_randomized_v1" + assert track.id == "atomic-task-randomized" + assert track.config["seeds"] == list(range(101, 117)) + assert [item["id"] for item in track.config["skills"]] == [ + "move_end_effector", + "move_joints", + "pick_up", + "move_held_object", + "place", + "press", + ] + assert all( + any(str(key).endswith(("_jitter_m", "_jitter_rad")) for key in item) + for item in track.config["skills"] + ) + + +def test_seeded_jitter_is_reproducible_bounded_and_stream_separated(): + kwargs = { + "amplitude": (0.1, 0.2, 0.0), + "seed": 107, + "dtype": torch.float64, + "device": torch.device("cpu"), + } + first = _seeded_jitter(stream=5, **kwargs) + repeated = _seeded_jitter(stream=5, **kwargs) + different_seed = _seeded_jitter(stream=5, **{**kwargs, "seed": 108}) + different_stream = _seeded_jitter(stream=6, **kwargs) + + assert torch.equal(first, repeated) + assert not torch.equal(first, different_seed) + assert not torch.equal(first, different_stream) + assert torch.all(torch.abs(first) <= torch.tensor((0.1, 0.2, 0.0))) + assert first.dtype == torch.float64 + + +def test_case_generation_seed_is_stable_and_stratified(): + value = _case_generation_seed(107, skill_index=2, case_index=3) + + assert value == _case_generation_seed(107, skill_index=2, case_index=3) + assert value != _case_generation_seed(108, skill_index=2, case_index=3) + assert value != _case_generation_seed(107, skill_index=3, case_index=3) + assert value != _case_generation_seed(107, skill_index=2, case_index=4) + + +def test_case_qpos_canonicalization_removes_sub_resolution_noise(): + first = torch.tensor([[0.123421, -0.456779]], dtype=torch.float32) + second = torch.tensor([[0.123419, -0.456781]], dtype=torch.float32) + + assert torch.equal(_canonical_case_qpos(first), _canonical_case_qpos(second)) + + +def test_randomization_manifest_records_contract_and_seed(): + parameters = _randomization_parameters( + { + "target_offset_jitter_m": [0.03, 0.02, 0.01], + "object_yaw_jitter_rad": 0.4, + "sample_count": 80, + }, + seed=113, + ) + + assert parameters == { + "enabled": True, + "seed": 113, + "distribution": "independent_uniform", + "ranges": { + "target_offset_jitter_m": [0.03, 0.02, 0.01], + "object_yaw_jitter_rad": 0.4, + }, + } + + +def test_franka_pgi_robot_provider_exposes_arm_hand_and_tcp(): + suite = load_suite("atomic_franka_pgi_curobo") + cfg = create_robot_provider(suite.robot).build_cfg() + + assert cfg.uid == "benchmark_franka_pgi" + assert len(cfg.control_parts["arm"]) == 7 + assert cfg.control_parts["hand"] == ["gripper_finger1_joint_1"] + assert cfg.solver_cfg["arm"].end_link_name == "fr3_link8" + assert cfg.solver_cfg["arm"].tcp[2][3] == pytest.approx(0.15) + assert len(cfg.init_qpos) == 9 + + +def test_atomic_invocation_pins_motion_generator_and_selected_planner(): + provider = create_atomic_skill_provider("move_end_effector") + invocation = provider.build_invocation( + Mock(control_part="manipulator"), + _atomic_case(), + Mock(motion_policy_planner="curobo"), + ) + + assert invocation.motion_policy.strategy == "motion_gen" + assert invocation.motion_policy.planner == "curobo" + assert invocation.skill_id == "move_end_effector" + assert invocation.binding.manipulators == {"primary": "manipulator"} + + +def test_atomic_skill_and_object_extensions_are_registry_driven(): + assert atomic_skill_provider_names() == ( + "move_end_effector", + "move_held_object", + "move_joints", + "pick_up", + "place", + "press", + ) + assert atomic_object_kind_names() == ("cube", "mesh") + with pytest.raises(ValueError, match="Unknown atomic object kind"): + create_atomic_object(Mock(), {"id": "new_object", "kind": "not_registered"}) + + +def test_atomic_primary_success_and_execution_efficiency_aggregate(): + case = _atomic_case() + metadata = [ + PlannerMetadata( + algorithm_id="curobo", + algorithm_role=AlgorithmRole.PRIMARY_BASELINE, + adapter="curobo", + config_hash="abc", + capabilities=frozenset({"eef_waypoint", "atomic_action"}), + supported_robots=("franka_pgi",), + ) + ] + record = TrialRecord( + suite_version=case.suite_version, + track=case.track, + scenario_id=case.scenario_id, + case_id=case.case_id, + algorithm_id="curobo", + algorithm_role=AlgorithmRole.PRIMARY_BASELINE, + model_revision="curobo-v2", + planner_config_hash="abc", + seed=case.seed, + repeat=0, + batch_size=1, + waypoint_count=1, + path_shape=case.path_shape, + start_state_bin=case.start_state_bin, + phase=TrialPhase.MEASURED, + cost_time_ms=20.0, + robot_id=case.robot_id, + skill_id=case.skill_id, + task_difficulty=case.task_difficulty, + primary_success=case.primary_success, + execution_time_ms=30.0, + end_to_end_time_ms=50.0, + trajectory_duration_s=1.5, + trajectory_waypoints=80, + outcomes=(_atomic_outcome(),), + ) + + aggregates = aggregate_results([record], metadata, [case], measured_trials=1) + metrics = aggregates["success_and_metrics"][0] + performance = aggregates["time_and_memory"][0] + leaderboard = aggregates["leaderboard"][0] + + assert metrics["primary_success"] == "task_success" + assert metrics["success_rate"] == pytest.approx(1.0) + assert metrics["execution_success_rate"] == pytest.approx(1.0) + assert metrics["task_success_rate"] == pytest.approx(1.0) + assert performance["execution_time_ms"] == pytest.approx(30.0) + assert performance["end_to_end_time_ms"] == pytest.approx(50.0) + assert leaderboard["overall_success_rate"] == pytest.approx(1.0) + assert leaderboard["task_success_rate"] == pytest.approx(1.0) + + +def test_atomic_case_manifest_retains_robot_skill_object_and_parameters(tmp_path): + case = _atomic_case() + path = write_case_manifest(tmp_path / "case_manifest.json", [case]) + payload = json.loads(path.read_text(encoding="utf-8")) + serialized = payload["cases"][0] + + assert payload["case_schema_version"] == 2 + assert serialized["robot_id"] == "franka_pgi" + assert serialized["skill_id"] == "move_end_effector" + assert serialized["primary_success"] == "task_success" + assert serialized["case_parameters"]["sample_count"] == 80 + assert serialized["validity_evidence"]["method"] == "independent_sequential_ik" + + +def test_should_record_case_respects_enable_failure_and_limit(): + disabled = VideoRecordCfg() + enabled = VideoRecordCfg(enabled=True) + with_failed = VideoRecordCfg(enabled=True, record_failed=True) + limited = VideoRecordCfg(enabled=True, case_limit=1) + + assert should_record_case(disabled, 0, True) is False + assert should_record_case(enabled, 0, True) is True + assert should_record_case(enabled, 0, False) is False + assert should_record_case(with_failed, 0, False) is True + assert should_record_case(limited, 0, True) is True + assert should_record_case(limited, 1, True) is False + + +def test_build_video_path_sanitizes_case_id(tmp_path): + path = build_video_path( + tmp_path, + "curobo", + "pick_up", + "atomic-task:pick_up:cube_top_center:s11", + ) + + assert path.parent == tmp_path + assert path.name == "curobo_pick_up_atomic-task_pick_up_cube_top_center_s11.mp4" + assert tmp_path.is_dir() + + +def test_default_scenario_record_replay_is_noop(tmp_path): + provider = FreeSpaceScenario() + path = provider.record_replay( + None, + _atomic_case(), + None, + output_dir=tmp_path, + algorithm_id="curobo", + video=VideoRecordCfg(enabled=True), + ) + + assert path is None + assert isinstance(provider, ScenarioProvider) + + +def test_atomic_record_replay_without_runtime_returns_none(tmp_path): + scenario = AtomicTaskScenario() + path = scenario.record_replay( + None, + _atomic_case(), + None, + output_dir=tmp_path, + algorithm_id="curobo", + video=VideoRecordCfg(enabled=True), + ) + assert path is None + + +def test_record_with_window_swallows_exceptions_and_does_not_raise(tmp_path): + sim = Mock() + sim.sim_config.width = 64 + sim.sim_config.height = 64 + sim.start_window_record.side_effect = RuntimeError("recorder failed") + sim.is_window_recording.return_value = False + + path = record_with_window( + sim, + VideoRecordCfg(enabled=True), + tmp_path / "failed.mp4", + lambda: None, + ) + + assert path is None + sim.wait_window_record_saves.assert_called() + + +def test_atomic_record_replay_swallows_recorder_errors(tmp_path): + scenario = AtomicTaskScenario() + sim = Mock() + sim.sim_config.width = 64 + sim.sim_config.height = 64 + sim.start_window_record.side_effect = RuntimeError("boom") + sim.is_window_recording.return_value = False + scenario.simulation = sim + scenario.robot = Mock() + scenario.reset_case = Mock() + + path = scenario.record_replay( + None, + _atomic_case(), + None, + output_dir=tmp_path, + algorithm_id="curobo", + video=VideoRecordCfg(enabled=True, record_failed=True), + ) + + assert path is None + scenario.reset_case.assert_called_once() + + +def test_atomic_failed_plan_records_static_hold(tmp_path): + scenario = AtomicTaskScenario() + sim = Mock() + sim.sim_config.width = 64 + sim.sim_config.height = 64 + sim.start_window_record.return_value = True + sim.is_window_recording.return_value = True + sim.stop_window_record.return_value = True + scenario.simulation = sim + scenario.robot = Mock() + scenario.track = Mock( + config={ + "physics": {"hold_steps": 2, "hold_sim_steps": 1, "steps_per_waypoint": 4} + } + ) + scenario.reset_case = Mock() + + path = scenario.record_replay( + None, + _atomic_case(), + None, + output_dir=tmp_path, + algorithm_id="curobo", + video=VideoRecordCfg(enabled=True, record_failed=True), + ) + + assert path is not None + assert path.name.startswith("curobo_move_end_effector_") + assert sim.update.call_count == 2 + + +def test_video_cfg_from_args_and_summary_notes(): + args = argparse.Namespace( + record_video=True, + record_failed_video=False, + video_case_limit=0, + video_fps=20, + video_width=640, + video_height=480, + video_max_memory=2048, + video_dir=None, + ) + cfg = video_cfg_from_args(args) + notes = summarize_video_recording(cfg, ()) + + assert cfg.enabled is True + assert cfg.record_failed is False + assert cfg.output_dir is None + assert "Video policy: disabled." not in notes + assert "videos=0" in notes + assert summarize_video_recording(VideoRecordCfg(), ()) == [ + "Video policy: disabled." + ] + with pytest.raises(ValueError, match="case_limit"): + VideoRecordCfg(enabled=True, case_limit=-1) + + +def test_runner_skips_video_outside_measured_phase(tmp_path): + suite = load_suite("atomic_franka_pgi_curobo") + specs = [spec for spec in suite.planners if spec.enabled] + runner = BenchmarkRunner( + suite, + specs, + device="cpu", + output_root=tmp_path, + video=VideoRecordCfg(enabled=True, record_failed=True), + ) + runner._run_dir = tmp_path + provider = Mock() + provider.record_replay.return_value = tmp_path / "should_not_write.mp4" + + path = runner._maybe_record_replay( + provider, + None, + _atomic_case(), + None, + "curobo", + TrialPhase.WARMUP, + ) + + assert path is None + provider.record_replay.assert_not_called() diff --git a/tests/benchmark/motion_generation/test_motion_generation_benchmark.py b/tests/benchmark/motion_generation/test_motion_generation_benchmark.py index d29a50af1..4425f4a08 100644 --- a/tests/benchmark/motion_generation/test_motion_generation_benchmark.py +++ b/tests/benchmark/motion_generation/test_motion_generation_benchmark.py @@ -24,6 +24,13 @@ import pytest import torch +from embodichain.lab.sim.atomic_actions import ( + AtomicActionEngine, + ControlPartCommandProfile, + GraspGoal, + HeldObjectPoseGoal, + MoveHeldObjectOptions, +) from embodichain.lab.sim.planners.curobo.curobo_planner import CuroboPlanner from embodichain.lab.sim.planners.utils import MoveType, PlanResult from scripts.benchmark.motion_generation.aggregation import aggregate_results @@ -54,6 +61,10 @@ from scripts.benchmark.motion_generation.run_benchmark import ( _apply_overrides, ) +from scripts.benchmark.motion_generation.scenarios.atomic_task import ( + AtomicTaskScenario, + create_atomic_skill_provider, +) def _translated_pose(x: float) -> torch.Tensor: @@ -114,6 +125,48 @@ def test_ordered_waypoint_requires_position_and_rotation_at_same_sample(): assert result["arrival_indices"] == [] +def test_press_waypoint_validation_accepts_half_turn_about_z_symmetry(): + target = torch.eye(4) + target[:3, 0] = -target[:3, 0] + target[:3, 1] = -target[:3, 1] + positions = torch.zeros(1, 2, 7) + strict_case = replace( + _case(), + scenario_id="press", + skill_id="press", + target_waypoints=target.reshape(1, 1, 4, 4), + ) + symmetric_case = replace( + strict_case, + case_parameters={"waypoint_rotation_symmetry": "half_turn_about_z"}, + ) + + strict = compute_case_outcomes( + _timed_plan_result(positions, success=True), + strict_case, + _MetricRobot(), + "arm", + validation_samples=8, + position_threshold_m=1.0e-4, + rotation_threshold_rad=1.0e-4, + joint_limit_tolerance_rad=1.0e-5, + ) + symmetric = compute_case_outcomes( + _timed_plan_result(positions, success=True), + symmetric_case, + _MetricRobot(), + "arm", + validation_samples=8, + position_threshold_m=1.0e-4, + rotation_threshold_rad=1.0e-4, + joint_limit_tolerance_rad=1.0e-5, + ) + + assert strict[0].failure_code == "waypoint_miss" + assert symmetric[0].motion_valid is True + assert symmetric[0].final_rotation_err_deg == pytest.approx(0.0) + + def test_waypoint_errors_use_threshold_greedy_arrivals(): """Continuous errors must come from the same matching as motion_valid.""" waypoints = torch.stack([_translated_pose(0.0), _translated_pose(0.10)]) @@ -315,6 +368,242 @@ def test_nmg_precision_and_external_accuracy_are_independently_configurable(): assert nmg.config["rot_eps"] == pytest.approx(0.20) +def test_nmg_model_revision_is_derived_from_runtime_model_path(): + from scripts.benchmark.motion_generation.config import PlannerSpecCfg + from scripts.benchmark.motion_generation.planners.base import PlannerContext + from scripts.benchmark.motion_generation.planners.nmg_onnx import NmgOnnxAdapter + + adapter = NmgOnnxAdapter( + PlannerSpecCfg( + id="nmg", + adapter="nmg_onnx", + role="candidate", + config={"onnx_model_path": "/models/unified-k3.onnx"}, + ), + PlannerContext( + robot=Mock(), + control_part="arm", + device=torch.device("cpu"), + sample_interval=1, + ), + ) + + assert adapter.metadata.model_revision == "unified-k3" + + +def test_seed_override_applies_to_atomic_tracks(): + suite = load_suite("atomic_franka_pgi_curobo_randomized") + + _apply_overrides(suite, seeds=[104]) + + assert suite.free_space.seeds == [104] + assert suite.enabled_tracks()[0].config["seeds"] == [104] + + +def test_atomic_task_antipodal_grasp_uses_standalone_generator(monkeypatch): + from scripts.tutorials.atomic_action import tutorial_utils + + vertices = torch.tensor( + [[-0.025, -0.025, -0.025], [0.025, 0.025, 0.025]], + dtype=torch.float32, + ) + triangles = torch.tensor([[0, 1, 0]], dtype=torch.int64) + entity = Mock(uid="atomic_cube") + entity.get_vertices.return_value = [vertices] + entity.get_triangles.return_value = [triangles] + handle = Mock(object_id="cube", entity=entity) + + candidate = torch.eye(4) + candidate[1, 1] = -1.0 + candidate[2, 2] = -1.0 + generator = Mock() + generator.get_valid_grasp_poses.return_value = [ + (candidate.unsqueeze(0), torch.tensor([0.25])) + ] + generator_factory = Mock(return_value=generator) + monkeypatch.setattr( + tutorial_utils, + "create_parallel_jaw_grasp_pose_generator", + generator_factory, + ) + + scenario = AtomicTaskScenario() + scenario.robot = Mock(device=torch.device("cpu")) + scenario.robot.compute_ik.return_value = ( + torch.tensor([True]), + torch.zeros(1, 7), + ) + result = scenario.resolve_antipodal_grasp( + handle, + torch.eye(4).unsqueeze(0), + torch.tensor([0.0, 0.0, -1.0]), + seed=11, + start_qpos=torch.zeros(1, 7), + pre_grasp_distance=0.15, + lift_height=0.16, + n_sample=321, + max_candidates=8, + alignment_max_angle_deg=10.0, + ) + + generator_factory.assert_called_once_with(n_sample=321, force_refresh=False) + call = generator.get_valid_grasp_poses.call_args.kwargs + assert torch.equal(call["mesh_vertices"], vertices) + assert torch.equal(call["mesh_triangles"], triangles) + assert torch.equal(result, candidate.unsqueeze(0)) + + +def test_atomic_task_pickup_builds_engine_owned_endpoint_binding(): + robot = Mock() + robot.device = torch.device("cpu") + robot.dof = 8 + robot.control_parts = {"arm": object(), "hand": object()} + robot.get_qpos.return_value = torch.zeros(1, 8) + robot.get_qvel.return_value = torch.zeros(1, 8) + robot.get_joint_ids.side_effect = lambda name: ( + list(range(7)) if name == "arm" else [7] + ) + motion_generator = Mock(robot=robot, device=torch.device("cpu")) + motion_generator.planner.cfg.planner_type = "stub" + engine = AtomicActionEngine( + motion_generator, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=torch.zeros(1), + grasp=torch.ones(1), + ) + }, + ) + entity = Mock(uid="atomic_cube") + scenario = AtomicTaskScenario() + scenario.robot = robot + scenario.control_part = "arm" + scenario.end_effector_part = "hand" + scenario._engine = engine + scenario._objects["cube"] = Mock(object_id="cube", entity=entity) + poses = torch.eye(4).reshape(1, 1, 4, 4).repeat(1, 3, 1, 1) + case = BenchmarkCase( + suite_version="test_v1", + track="atomic-task", + scenario_id="pick_up", + case_id="atomic-task:pick_up:cube:s11", + seed=11, + batch_size=1, + num_waypoints=3, + path_shape="approach_grasp_lift", + start_state_bin="pre_pick", + start_qpos=torch.zeros(1, 7), + target_waypoints=poses, + reference_qpos=torch.zeros(1, 3, 7), + skill_id="pick_up", + object_id="cube", + case_parameters={ + "sample_count": 12, + "approach_direction": [0.0, 0.0, -1.0], + "pre_grasp_distance_m": 0.15, + "lift_height_m": 0.16, + "hand_interp_steps": 4, + }, + ) + + invocation = create_atomic_skill_provider("pick_up").build_invocation( + scenario, + case, + Mock(motion_policy_planner="curobo"), + ) + + assert isinstance(invocation.goal, GraspGoal) + assert invocation.goal.semantics.entity_id == "atomic_cube" + assert invocation.binding.owner_id == engine.binding_owner_id + assert set(invocation.binding.endpoint_keys) == { + ("primary", "motion"), + ("primary", "grasp"), + } + + +def test_atomic_task_prepare_planner_registers_benchmark_objects_in_scene(): + robot = Mock() + robot.device = torch.device("cpu") + robot.dof = 7 + robot.control_parts = {"arm": object()} + robot.get_qpos.return_value = torch.zeros(1, 7) + robot.get_qvel.return_value = torch.zeros(1, 7) + motion_generator = Mock(robot=robot, device=torch.device("cpu")) + motion_generator.planner.cfg.planner_type = "stub" + adapter = Mock() + adapter.require_motion_generator.return_value = motion_generator + + object_pose = torch.eye(4).unsqueeze(0) + entity = Mock(uid="atomic_cube") + entity.get_local_pose.return_value = object_pose + scenario = AtomicTaskScenario() + scenario._objects["cube"] = Mock(object_id="cube", entity=entity) + + scenario.prepare_planner(adapter, Mock()) + + scene = scenario.require_engine().initial_context().scene + assert tuple(scene.entities) == ("atomic_cube",) + assert torch.equal(scene.entities["atomic_cube"].pose, object_pose) + + +def test_atomic_task_move_held_object_uses_current_options_contract(): + robot = Mock() + robot.device = torch.device("cpu") + robot.dof = 8 + robot.control_parts = {"arm": object(), "hand": object()} + robot.get_qpos.return_value = torch.zeros(1, 8) + robot.get_qvel.return_value = torch.zeros(1, 8) + robot.get_joint_ids.side_effect = lambda name: ( + list(range(7)) if name == "arm" else [7] + ) + motion_generator = Mock(robot=robot, device=torch.device("cpu")) + motion_generator.planner.cfg.planner_type = "stub" + engine = AtomicActionEngine( + motion_generator, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=torch.zeros(1), + grasp=torch.ones(1), + ) + }, + ) + scenario = AtomicTaskScenario() + scenario.robot = robot + scenario.control_part = "arm" + scenario.end_effector_part = "hand" + scenario._engine = engine + target_object_pose = torch.eye(4).unsqueeze(0) + case = BenchmarkCase( + suite_version="test_v1", + track="atomic-task", + scenario_id="move_held_object", + case_id="atomic-task:move_held_object:cube:s11", + seed=11, + batch_size=1, + num_waypoints=1, + path_shape="held_object_transport", + start_state_bin="object_held", + start_qpos=torch.zeros(1, 7), + target_waypoints=target_object_pose.unsqueeze(1), + reference_qpos=torch.zeros(1, 1, 7), + skill_id="move_held_object", + object_id="cube", + case_parameters={ + "sample_count": 12, + "target_object_pose": target_object_pose.tolist(), + }, + ) + + invocation = create_atomic_skill_provider("move_held_object").build_invocation( + scenario, + case, + Mock(), + ) + + assert isinstance(invocation.goal, HeldObjectPoseGoal) + assert type(invocation.skill_options) is MoveHeldObjectOptions + + @pytest.mark.parametrize("override", [{"nmg_pos_eps": 0.0}, {"nmg_rot_eps": -0.1}]) def test_nmg_precision_rejects_non_positive_values(override): suite = load_suite("smoke") diff --git a/tests/sim/atomic_actions/test_engine.py b/tests/sim/atomic_actions/test_engine.py index 0622cc7f5..b1fde477e 100644 --- a/tests/sim/atomic_actions/test_engine.py +++ b/tests/sim/atomic_actions/test_engine.py @@ -334,6 +334,21 @@ def test_engine_compile_holds_failed_rows_for_remaining_actions() -> None: assert torch.all(compiled.trajectory.positions[1] == 0.0) +def test_engine_compile_preserves_opted_in_failed_trajectory_for_diagnostics() -> None: + engine = _engine() + engine.motion_generator.planner.preserve_failed_plan_positions = True + engine.register(StubAction()) + target = torch.tensor([[1.0, 1.0, 1.0], [float("nan"), 2.0, 2.0]]) + + compiled = engine.compile((_invocation(engine, target),)) + + assert compiled.plan_success.tolist() == [True, False] + assert torch.equal( + compiled.trajectory.positions[1, -1], torch.tensor([0.0, 2.0, 2.0]) + ) + assert torch.all(compiled.projected_context.robot.qpos[1] == 0.0) + + def test_engine_compile_empty_sequence_is_successful_noop() -> None: engine = _engine() context = engine.initial_context() diff --git a/tests/sim/planners/test_motion_generator_batched.py b/tests/sim/planners/test_motion_generator_batched.py index 4b84350d9..9887c4c43 100644 --- a/tests/sim/planners/test_motion_generator_batched.py +++ b/tests/sim/planners/test_motion_generator_batched.py @@ -162,6 +162,56 @@ def test_direct_cartesian_planner_requires_joint_fallback_inputs(): ) +def test_motion_generator_dispatches_heterogeneous_waypoints_by_capability(): + planner = Mock() + planner.supports_heterogeneous_waypoints = True + planner.supports_move_type.side_effect = lambda move_type: move_type in { + MoveType.EEF_MOVE, + MoveType.JOINT_MOVE, + } + planner.preserve_plan_samples = True + planner.default_plan_options.return_value = PlanOptions() + planner.with_motion_context.side_effect = ( + lambda options, *, start_qpos, control_part: options + ) + planner.plan.return_value = PlanResult( + success=torch.ones(1, dtype=torch.bool), + positions=torch.zeros(1, 5, 2), + dt=torch.full((1, 5), 0.01), + ) + generator = object.__new__(MotionGenerator) + generator.planner = planner + generator.device = torch.device("cpu") + targets = [ + PlanState.from_xpos(torch.eye(4).unsqueeze(0)), + PlanState.from_qpos(torch.zeros(1, 2)), + ] + + result = generator.generate( + targets, + MotionGenOptions(start_qpos=torch.zeros(1, 2), control_part="arm"), + ) + + assert result.success.all().item() + assert planner.plan.call_args.kwargs["target_states"] is targets + + +def test_motion_generator_rejects_heterogeneous_waypoints_without_capability(): + planner = _DirectCartesianPlanner() + generator = object.__new__(MotionGenerator) + generator.planner = planner + generator.device = torch.device("cpu") + + with pytest.raises(ValueError, match="does not support heterogeneous"): + generator.generate( + [ + PlanState.from_xpos(torch.eye(4).unsqueeze(0)), + PlanState.from_qpos(torch.zeros(1, 2)), + ], + MotionGenOptions(start_qpos=torch.zeros(1, 2), control_part="arm"), + ) + + def test_bind_collision_world_copies_caller_options() -> None: planner = Mock() planner.collision_world_info = _collision_world_info(("obstacle",)) diff --git a/tests/sim/planners/test_neural_batched.py b/tests/sim/planners/test_neural_batched.py index 78055af88..6525ab9ee 100644 --- a/tests/sim/planners/test_neural_batched.py +++ b/tests/sim/planners/test_neural_batched.py @@ -16,7 +16,11 @@ from __future__ import annotations import torch -import pytest + + +def _set_identity_policy_frames(planner) -> None: + planner._policy_frame_from_world = torch.eye(4) + planner._runtime_tcp_from_policy_tcp = torch.eye(4) class TestNeuralParseWaypoints: @@ -28,6 +32,9 @@ def test_parse_waypoints_batched(self): planner = NeuralPlanner.__new__(NeuralPlanner) planner.device = torch.device("cpu") planner._num_waypoints = 4 + planner._action_dim = 7 + planner._intermediate_orientation = True + _set_identity_policy_frames(planner) B = 3 states = [ PlanState.from_xpos( @@ -36,10 +43,16 @@ def test_parse_waypoints_batched(self): ) for _ in range(2) ] - pos, quat, mask, k = planner._parse_waypoints(states) + pos, quat, joint, mask, pos_mask, rot_mask, joint_mask, k = ( + planner._parse_waypoints(states) + ) assert pos.shape == (B, 4, 3) assert quat.shape == (B, 4, 4) + assert joint.shape == (B, 4, 7) assert mask.shape == (B, 4) + assert torch.equal(pos_mask, mask) + assert torch.equal(rot_mask, mask) + assert torch.count_nonzero(joint_mask) == 0 assert k == 2 @@ -58,9 +71,11 @@ def test_plan_returns_batched_success(self, monkeypatch): planner._max_steps = 5 planner._pos_eps = 1e9 # always reached planner._rot_eps = 1e9 + planner._joint_eps = 1e9 planner._intermediate_orientation = True planner._use_relative_obs = False - planner._obs_dim = 57 # 7+7+4*3+4*4+4+4+7 + planner._obs_dim = 101 + _set_identity_policy_frames(planner) planner.cfg = type( "c", (), @@ -72,9 +87,8 @@ def test_plan_returns_batched_success(self, monkeypatch): }, )() - # stub actor: returns zeros so qpos never changes but eps is huge -> reached - planner._actor = lambda obs: torch.zeros(obs.shape[0], 7) - planner._normalizer = type("n", (), {"normalize": lambda self, o: o})() + # Stub ONNX policy: qpos stays fixed but the huge eps marks each goal reached. + planner._policy = lambda obs: torch.zeros(obs.shape[0], 7) # stub robot FK + limits class _Robot: @@ -125,10 +139,11 @@ def test_converged_env_holds_qpos(self): planner._max_steps = 5 planner._pos_eps = 1e9 planner._rot_eps = 1e9 + planner._joint_eps = 1e9 planner._intermediate_orientation = True planner._use_relative_obs = False - # joint(7)+ee(7)+wp_pos(4*3)+wp_quat(4*4)+onehot(4)+valid(4)+last_act(7)=57 - planner._obs_dim = 57 + planner._obs_dim = 101 + _set_identity_policy_frames(planner) planner.cfg = type( "c", (), @@ -140,9 +155,8 @@ def test_converged_env_holds_qpos(self): }, )() - # actor: non-trivial action so qpos would drift if not masked - planner._actor = lambda obs: torch.ones(obs.shape[0], 7) * 0.5 - planner._normalizer = type("n", (), {"normalize": lambda self, o: o})() + # Non-trivial policy action so qpos would drift if not masked. + planner._policy = lambda obs: torch.ones(obs.shape[0], 7) * 0.5 class _Robot: num_instances = 3 @@ -168,7 +182,17 @@ def compute_fk(self, qpos=None, name=None, to_matrix=True): # episode_k = 2 (2 waypoints). Env 0: reached=True every step -> # active_idx 0->1 (step0), 1->2 (step1) => converges after step 1. # Envs 1,2: reached=False always -> never converge within max_steps. - def fake_reach(ee_pose, wp_pos, wp_quat, active_idx, episode_k): + def fake_reach( + joint_pos, + ee_pose, + wp_pos, + wp_quat, + wp_joint, + pos_mask, + rot_mask, + joint_mask, + active_idx, + ): r = torch.zeros(ee_pose.shape[0], dtype=torch.bool, device=ee_pose.device) r[0] = True return r diff --git a/tests/sim/planners/test_neural_planner.py b/tests/sim/planners/test_neural_planner.py index c1e2d2786..e7e878ee8 100644 --- a/tests/sim/planners/test_neural_planner.py +++ b/tests/sim/planners/test_neural_planner.py @@ -27,52 +27,38 @@ NeuralPlannerCfg, PlanState, ) -from embodichain.lab.sim.planners.neural_planner import ( - NeuralPlanOptions, - _WaypointTransformerActor, -) +from embodichain.lab.sim.planners.neural_planner import NeuralPlanOptions +from embodichain.lab.sim.planners import neural_planner as neural_planner_module from embodichain.lab.sim.sim_manager import SimulationManager NUM_ARM_JOINTS = 7 -NUM_WAYPOINTS = 3 -OBS_DIM = 28 + 9 * NUM_WAYPOINTS -HIDDEN_DIM = 32 - - -def _create_fake_checkpoint(tmp_path) -> str: - actor = _WaypointTransformerActor( - obs_dim=OBS_DIM, - action_dim=NUM_ARM_JOINTS, - num_waypoints=NUM_WAYPOINTS, - use_relative_obs=True, - hidden_dim=HIDDEN_DIM, - transformer_nhead=4, - transformer_num_layers=1, - ) - checkpoint = { - "agent": {f"actor_mean.{k}": v for k, v in actor.state_dict().items()}, - "obs_normalizer": { - "mean": torch.zeros(OBS_DIM), - "var": torch.ones(OBS_DIM), - "count": 1.0, - }, - "args": { - "policy_arch": "transformer", - "hidden_dim": HIDDEN_DIM, - "transformer_nhead": 4, - "transformer_num_layers": 1, - "transformer_ff_dim": 0, - "waypoint_max": NUM_WAYPOINTS, - "waypoint_use_relative_obs": True, - "waypoint_intermediate_orientation": True, - "max_episode_steps": 3, - "waypoint_pos_threshold": 0.05, - "waypoint_rot_threshold": 0.3, - }, - } - checkpoint_path = tmp_path / "fake_neural_planner.pt" - torch.save(checkpoint, checkpoint_path) - return str(checkpoint_path) +NUM_WAYPOINTS = 8 +OBS_DIM = 300 + + +def _create_fake_onnx_model(tmp_path) -> str: + model_path = tmp_path / "fake_neural_planner.onnx" + model_path.write_bytes(b"fake-onnx") + return str(model_path) + + +class FakeOnnxPolicy: + obs_dim = OBS_DIM + fixed_batch_size = 1 + + def __init__(self, path, providers=None): + self.path = path + self.providers = providers + self.last_obs = None + + def __call__(self, obs: torch.Tensor) -> torch.Tensor: + self.last_obs = obs.clone() + return torch.zeros(obs.shape[0], NUM_ARM_JOINTS, device=obs.device) + + +@pytest.fixture(autouse=True) +def _mock_onnx_policy(monkeypatch): + monkeypatch.setattr(neural_planner_module, "_OnnxPolicy", FakeOnnxPolicy) class FakeRobot: @@ -113,10 +99,15 @@ def get_robot(self, uid: str) -> FakeRobot: def test_neural_planner_is_registered(): assert MotionGenerator._support_planner_dict["neural"][0] is NeuralPlanner assert MotionGenerator._support_planner_dict["neural"][1] is NeuralPlannerCfg + assert NeuralPlanner.preserve_plan_samples is True + assert NeuralPlanner.preserve_failed_plan_positions is True + assert NeuralPlanner.supported_move_types == frozenset( + {MoveType.EEF_MOVE, MoveType.JOINT_MOVE} + ) -def test_neural_planner_generate_with_fake_checkpoint(tmp_path, monkeypatch): - checkpoint_path = _create_fake_checkpoint(tmp_path) +def test_neural_planner_generate_with_fake_onnx_model(tmp_path, monkeypatch): + model_path = _create_fake_onnx_model(tmp_path) fake_sim = FakeSimulationManager() monkeypatch.setattr( SimulationManager, "get_instance", classmethod(lambda cls: fake_sim) @@ -126,7 +117,7 @@ def test_neural_planner_generate_with_fake_checkpoint(tmp_path, monkeypatch): cfg=MotionGenCfg( planner_cfg=NeuralPlannerCfg( robot_uid="fake_robot", - checkpoint_path=checkpoint_path, + onnx_model_path=model_path, control_part="main_arm", ) ) @@ -152,7 +143,7 @@ def test_neural_planner_generate_with_fake_checkpoint(tmp_path, monkeypatch): def test_neural_planner_uses_plan_opts_start_qpos(tmp_path, monkeypatch): - checkpoint_path = _create_fake_checkpoint(tmp_path) + model_path = _create_fake_onnx_model(tmp_path) fake_sim = FakeSimulationManager() monkeypatch.setattr( SimulationManager, "get_instance", classmethod(lambda cls: fake_sim) @@ -162,7 +153,7 @@ def test_neural_planner_uses_plan_opts_start_qpos(tmp_path, monkeypatch): cfg=MotionGenCfg( planner_cfg=NeuralPlannerCfg( robot_uid="fake_robot", - checkpoint_path=checkpoint_path, + onnx_model_path=model_path, control_part="main_arm", ) ) @@ -185,7 +176,7 @@ def test_neural_planner_uses_plan_opts_start_qpos(tmp_path, monkeypatch): def test_neural_planner_rejects_short_start_qpos(tmp_path, monkeypatch): - checkpoint_path = _create_fake_checkpoint(tmp_path) + model_path = _create_fake_onnx_model(tmp_path) fake_sim = FakeSimulationManager() monkeypatch.setattr( SimulationManager, "get_instance", classmethod(lambda cls: fake_sim) @@ -195,7 +186,7 @@ def test_neural_planner_rejects_short_start_qpos(tmp_path, monkeypatch): cfg=MotionGenCfg( planner_cfg=NeuralPlannerCfg( robot_uid="fake_robot", - checkpoint_path=checkpoint_path, + onnx_model_path=model_path, control_part="main_arm", ) ) @@ -216,7 +207,7 @@ def test_neural_planner_rejects_short_start_qpos(tmp_path, monkeypatch): def test_neural_planner_returns_velocities_and_accelerations(tmp_path, monkeypatch): - checkpoint_path = _create_fake_checkpoint(tmp_path) + model_path = _create_fake_onnx_model(tmp_path) fake_sim = FakeSimulationManager() monkeypatch.setattr( SimulationManager, "get_instance", classmethod(lambda cls: fake_sim) @@ -226,7 +217,7 @@ def test_neural_planner_returns_velocities_and_accelerations(tmp_path, monkeypat cfg=MotionGenCfg( planner_cfg=NeuralPlannerCfg( robot_uid="fake_robot", - checkpoint_path=checkpoint_path, + onnx_model_path=model_path, control_part="main_arm", ) ) @@ -252,6 +243,243 @@ def test_neural_planner_returns_velocities_and_accelerations(tmp_path, monkeypat assert torch.isfinite(result.accelerations).all() +def test_neural_planner_disables_grad_for_all_fk_calls(tmp_path, monkeypatch): + model_path = _create_fake_onnx_model(tmp_path) + fake_sim = FakeSimulationManager() + grad_states = [] + original_compute_fk = fake_sim.robot.compute_fk + + def checked_compute_fk(*args, **kwargs): + grad_states.append(torch.is_grad_enabled()) + return original_compute_fk(*args, **kwargs) + + monkeypatch.setattr(fake_sim.robot, "compute_fk", checked_compute_fk) + monkeypatch.setattr( + SimulationManager, "get_instance", classmethod(lambda cls: fake_sim) + ) + planner = NeuralPlanner( + NeuralPlannerCfg( + robot_uid="fake_robot", + onnx_model_path=model_path, + control_part="main_arm", + ) + ) + + planner.plan( + [PlanState.single(move_type=MoveType.EEF_MOVE, xpos=torch.eye(4))], + NeuralPlanOptions( + control_part="main_arm", + start_qpos=torch.zeros(NUM_ARM_JOINTS), + max_steps=1, + ), + ) + + assert grad_states + assert not any(grad_states) + + +def test_neural_planner_builds_unified_300d_cartesian_observation( + tmp_path, monkeypatch +): + model_path = _create_fake_onnx_model(tmp_path) + fake_sim = FakeSimulationManager() + monkeypatch.setattr( + SimulationManager, "get_instance", classmethod(lambda cls: fake_sim) + ) + planner = NeuralPlanner( + NeuralPlannerCfg( + robot_uid="fake_robot", + onnx_model_path=model_path, + control_part="main_arm", + ) + ) + joint = torch.zeros(1, 7) + eef = torch.tensor([[0.1, 0.2, 0.3, 0.0, 0.0, 0.0, 1.0]]) + waypoint_pos = torch.zeros(1, 8, 3) + waypoint_quat = torch.zeros(1, 8, 4) + waypoint_quat[..., 3] = 1.0 + waypoint_joint = torch.zeros(1, 8, 7) + valid = torch.zeros(1, 8) + valid[:, :2] = 1.0 + pos_mask = valid.clone() + rot_mask = valid.clone() + joint_mask = torch.zeros_like(valid) + obs = planner._build_obs( + joint, + eef, + waypoint_pos, + waypoint_quat, + waypoint_joint, + valid, + pos_mask, + rot_mask, + joint_mask, + torch.zeros(1, dtype=torch.long), + torch.zeros(1, 7), + ) + + assert obs.shape == (1, 300) + # Unified layout semantic blocks: active, valid, pos, rot, joint masks. + semantic_start = 7 + 7 + 8 * (3 + 4 + 7) + expected_active = torch.zeros_like(valid) + expected_active[:, 0] = 1.0 + assert torch.equal(obs[:, semantic_start : semantic_start + 8], expected_active) + assert torch.equal(obs[:, semantic_start + 8 : semantic_start + 16], valid) + assert torch.equal(obs[:, semantic_start + 16 : semantic_start + 24], valid) + assert torch.equal(obs[:, semantic_start + 24 : semantic_start + 32], valid) + assert torch.count_nonzero(obs[:, semantic_start + 32 : semantic_start + 40]) == 0 + + +def test_neural_planner_builds_joint_constraint_observation(tmp_path, monkeypatch): + model_path = _create_fake_onnx_model(tmp_path) + fake_sim = FakeSimulationManager() + monkeypatch.setattr( + SimulationManager, "get_instance", classmethod(lambda cls: fake_sim) + ) + planner = NeuralPlanner( + NeuralPlannerCfg( + robot_uid="fake_robot", + onnx_model_path=model_path, + control_part="main_arm", + ) + ) + target = torch.tensor([[0.1, -0.2, 0.3, 0.0, 0.2, -0.1, 0.4]]) + parsed = planner._parse_waypoints( + [PlanState.from_qpos(target, move_type=MoveType.JOINT_MOVE)] + ) + ( + waypoint_pos, + waypoint_quat, + waypoint_joint, + valid, + pos_mask, + rot_mask, + joint_mask, + _, + ) = parsed + obs = planner._build_obs( + torch.zeros(1, 7), + torch.tensor([[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]]), + waypoint_pos, + waypoint_quat, + waypoint_joint, + valid, + pos_mask, + rot_mask, + joint_mask, + torch.zeros(1, dtype=torch.long), + torch.zeros(1, 7), + ) + + semantic_start = 7 + 7 + 8 * (3 + 4 + 7) + relative_start = semantic_start + 5 * 8 + 7 + waypoint_type_start = relative_start + 7 + 8 * (3 + 4 + 7) + assert torch.equal(joint_mask[:, 0], torch.ones(1)) + assert torch.count_nonzero(pos_mask) == 0 + assert torch.count_nonzero(rot_mask) == 0 + assert torch.allclose(obs[:, relative_start : relative_start + 7], target) + assert obs[0, waypoint_type_start].item() == pytest.approx(2.0) + + +def test_neural_planner_parses_ordered_pose_then_joint_sequence(tmp_path, monkeypatch): + model_path = _create_fake_onnx_model(tmp_path) + fake_sim = FakeSimulationManager() + monkeypatch.setattr( + SimulationManager, "get_instance", classmethod(lambda cls: fake_sim) + ) + planner = NeuralPlanner( + NeuralPlannerCfg( + robot_uid="fake_robot", + onnx_model_path=model_path, + control_part="main_arm", + ) + ) + pose = torch.eye(4).unsqueeze(0) + joint = torch.tensor([[0.1, -0.2, 0.3, 0.0, 0.2, -0.1, 0.4]]) + + ( + _, + _, + parsed_joint, + valid, + pos_mask, + rot_mask, + joint_mask, + episode_k, + ) = planner._parse_waypoints( + [ + PlanState.from_xpos(pose, move_type=MoveType.EEF_MOVE), + PlanState.from_qpos(joint, move_type=MoveType.JOINT_MOVE), + ] + ) + + assert episode_k == 2 + assert torch.equal(valid[0, :2], torch.ones(2)) + assert torch.equal(pos_mask[0, :2], torch.tensor([1.0, 0.0])) + assert torch.equal(rot_mask[0, :2], torch.tensor([1.0, 0.0])) + assert torch.equal(joint_mask[0, :2], torch.tensor([0.0, 1.0])) + assert torch.allclose(parsed_joint[0, 1], joint[0]) + + +def test_neural_planner_accepts_joint_move_goal(tmp_path, monkeypatch): + model_path = _create_fake_onnx_model(tmp_path) + fake_sim = FakeSimulationManager() + monkeypatch.setattr( + SimulationManager, "get_instance", classmethod(lambda cls: fake_sim) + ) + planner = NeuralPlanner( + NeuralPlannerCfg( + robot_uid="fake_robot", + onnx_model_path=model_path, + control_part="main_arm", + ) + ) + + result = planner.plan( + [PlanState.from_qpos(torch.zeros(1, 7), move_type=MoveType.JOINT_MOVE)], + NeuralPlanOptions( + control_part="main_arm", + start_qpos=torch.zeros(NUM_ARM_JOINTS), + max_steps=1, + ), + ) + + assert result.success.all().item() + + +def test_neural_planner_applies_policy_frame_and_tcp_transforms(tmp_path, monkeypatch): + model_path = _create_fake_onnx_model(tmp_path) + fake_sim = FakeSimulationManager() + monkeypatch.setattr( + SimulationManager, "get_instance", classmethod(lambda cls: fake_sim) + ) + left = [ + [-1.0, 0.0, 0.0, 0.0], + [0.0, -1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 0.0, 1.0], + ] + right = [ + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, -0.0466], + [0.0, 0.0, 0.0, 1.0], + ] + planner = NeuralPlanner( + NeuralPlannerCfg( + robot_uid="fake_robot", + onnx_model_path=model_path, + policy_frame_from_world=left, + runtime_tcp_from_policy_tcp=right, + ) + ) + pose = torch.eye(4).unsqueeze(0) + pose[:, :3, 3] = torch.tensor([[-0.5, 0.1, 0.4]]) + transformed = planner._to_policy_frame(pose) + + assert torch.allclose(transformed[0, :3, 3], torch.tensor([0.5, -0.1, 0.3534])) + + def test_neural_planner_finite_diff_helper(): """Finite-difference estimates match a known polynomial trajectory.""" b, n, dof = 2, 5, 7 @@ -282,7 +510,7 @@ def test_neural_planner_finite_diff_helper(): def test_motion_generator_neural_propagates_motion_gen_options(tmp_path, monkeypatch): - checkpoint_path = _create_fake_checkpoint(tmp_path) + model_path = _create_fake_onnx_model(tmp_path) fake_sim = FakeSimulationManager() monkeypatch.setattr( SimulationManager, "get_instance", classmethod(lambda cls: fake_sim) @@ -292,7 +520,7 @@ def test_motion_generator_neural_propagates_motion_gen_options(tmp_path, monkeyp cfg=MotionGenCfg( planner_cfg=NeuralPlannerCfg( robot_uid="fake_robot", - checkpoint_path=checkpoint_path, + onnx_model_path=model_path, ) ) ) @@ -312,7 +540,7 @@ def test_motion_generator_neural_propagates_motion_gen_options(tmp_path, monkeyp def test_motion_generator_neural_preserves_native_eef_targets(tmp_path, monkeypatch): - checkpoint_path = _create_fake_checkpoint(tmp_path) + model_path = _create_fake_onnx_model(tmp_path) fake_sim = FakeSimulationManager() monkeypatch.setattr( SimulationManager, "get_instance", classmethod(lambda cls: fake_sim) @@ -322,7 +550,7 @@ def test_motion_generator_neural_preserves_native_eef_targets(tmp_path, monkeypa cfg=MotionGenCfg( planner_cfg=NeuralPlannerCfg( robot_uid="fake_robot", - checkpoint_path=checkpoint_path, + onnx_model_path=model_path, control_part="main_arm", ) ) @@ -344,13 +572,61 @@ def test_motion_generator_neural_preserves_native_eef_targets(tmp_path, monkeypa assert options.is_interpolate is True -def test_safe_torch_load_roundtrip(tmp_path): - checkpoint = {"agent": torch.tensor([1.0, 2.0, 3.0])} - path = tmp_path / "checkpoint.pt" - torch.save(checkpoint, path) +def test_motion_generator_neural_preserves_failed_rollout_positions( + tmp_path, monkeypatch +): + class MovingFakeOnnxPolicy(FakeOnnxPolicy): + def __call__(self, obs: torch.Tensor) -> torch.Tensor: + self.last_obs = obs.clone() + action = torch.zeros(obs.shape[0], NUM_ARM_JOINTS, device=obs.device) + action[:, 0] = 1.0 + return action - from embodichain.lab.sim.planners.neural_planner import _safe_torch_load + monkeypatch.setattr(neural_planner_module, "_OnnxPolicy", MovingFakeOnnxPolicy) + model_path = _create_fake_onnx_model(tmp_path) + fake_sim = FakeSimulationManager() + monkeypatch.setattr( + SimulationManager, "get_instance", classmethod(lambda cls: fake_sim) + ) + + motion_generator = MotionGenerator( + cfg=MotionGenCfg( + planner_cfg=NeuralPlannerCfg( + robot_uid="fake_robot", + onnx_model_path=model_path, + control_part="main_arm", + max_steps=2, + ) + ) + ) + target = torch.eye(4) + target[0, 3] = 10.0 + + result = motion_generator.generate( + target_states=[PlanState.single(move_type=MoveType.EEF_MOVE, xpos=target)], + options=MotionGenOptions( + control_part="main_arm", + start_qpos=torch.zeros(NUM_ARM_JOINTS), + ), + ) + + assert not result.success.all().item() + assert result.positions is not None + assert result.positions[0, 1, 0] > result.positions[0, 0, 0] + + +def test_neural_planner_rejects_pytorch_checkpoint(tmp_path, monkeypatch): + pytorch_checkpoint_path = tmp_path / "checkpoint.pt" + pytorch_checkpoint_path.write_bytes(b"not-an-onnx-model") + fake_sim = FakeSimulationManager() + monkeypatch.setattr( + SimulationManager, "get_instance", classmethod(lambda cls: fake_sim) + ) - loaded = _safe_torch_load(path, map_location=torch.device("cpu")) - assert "agent" in loaded - assert torch.equal(loaded["agent"], checkpoint["agent"]) + with pytest.raises(ValueError, match="only accepts standalone .onnx"): + NeuralPlanner( + NeuralPlannerCfg( + robot_uid="fake_robot", + onnx_model_path=str(pytorch_checkpoint_path), + ) + )