Add env-batched parallel motion generation#349
Merged
Conversation
Env-batched (B=num_envs) parallel motion generation across the planner (BasePlanner), MotionGenerator facade, and atomic-action layer. Settles: batched PlanState/PlanResult, fork-multiprocessing TOPPRA fan-out, natively-batched NeuralPlanner, dual motion-source TrajectoryBuilder (ik_interp default + motion_gen opt-in), per-env failure propagation. Co-Authored-By: Claude <noreply@anthropic.com>
TDD task-by-task plan covering: batched PlanState/PlanResult, fork-mp TOPPRA fan-out, natively batched NeuralPlanner, batched MotionGenerator, TrajectoryBuilder motion_source strategy, per-env ActionResult/engine failure propagation, 6-action migration, caller/tutorial migration, integration tests with TOPPRA regression anchor. Co-Authored-By: Claude <noreply@anthropic.com>
Co-Authored-By: Claude <noreply@anthropic.com>
Co-Authored-By: Claude <noreply@anthropic.com>
Co-Authored-By: Claude <noreply@anthropic.com>
Co-Authored-By: Claude <noreply@anthropic.com>
…nstances guard Co-Authored-By: Claude <noreply@anthropic.com>
…ces guard Co-Authored-By: Claude <noreply@anthropic.com>
… holds Co-Authored-By: Claude <noreply@anthropic.com>
Co-Authored-By: Claude <noreply@anthropic.com>
…atched - Rewrite interpolate_trajectory to preserve the B dim across joint-space and Cartesian branches. - Add interpolate_xpos_batched helper in utils.py for batched pose interpolation. - Fix valid mask dtype: cast compute_batch_ik success to bool before boolean indexing. - Update single-env tests to PlanState.single(...) and assert (1, M, DOF) shapes. - Add TestInterpolateBatched for joint-space batched interpolation. Co-Authored-By: Claude <noreply@anthropic.com>
…hrough - MoveEndEffector.execute: per-env success tensor, cfg=self.cfg, remove scalar short-circuit - MoveJoints.execute: per-env ones success tensor - PickUp/Place/MoveHeldObject/Press: add if not ok.all().item() shim to prevent RuntimeError on tensor ok (scalar short-circuit preserved until G2/G3) - test_actions.py: update MoveEndEffector + MoveJoints assertions to result.success.all() + result.success.shape == (NUM_ENVS,) Co-Authored-By: Claude <noreply@anthropic.com>
- Tutorials: change is_success bool checks to is_success.all() for batched (B,) tensor return from AtomicActionEngine.run() - motion_generator tutorial: use PlanState.single() for B=1, index positions[0] for batched PlanResult - neural_planner benchmark: use PlanState.single(), update plan_ik_interpolate to return batched PlanResult, adapt _compute_result_metrics/_trajectory_fk_poses/_final_eef_pose for (B,N,DOF) positions and (B,) success/duration Co-Authored-By: Claude <noreply@anthropic.com>
Co-Authored-By: Claude <noreply@anthropic.com>
Co-Authored-By: Claude <noreply@anthropic.com>
Co-Authored-By: Claude <noreply@anthropic.com>
…l motion_gen, _fail tensor - engine.run: break on all-dead instead of padding (shape consistency) - migrate 3 tableware task .generate() callers to batched PlanState.single - _plan_motion_gen: is_interpolate=False for neural planner (EEF_MOVE compat) - remove stale build_motion_gen_options helper - _fail helpers return (B,) all-False tensor Co-Authored-By: Claude <noreply@anthropic.com>
Comment on lines
+70
to
+81
| b = bs.pop() if bs else 1 | ||
| if expected_b is not None and b != expected_b: | ||
| logger.log_error( | ||
| f"Batch dim B={b} does not match robot.num_instances={expected_b}", | ||
| ValueError, | ||
| ) | ||
| if robot_num_instances is not None and b not in (1, robot_num_instances): | ||
| logger.log_error( | ||
| f"Batch dim B={b} must be 1 or robot.num_instances={robot_num_instances}", | ||
| ValueError, | ||
| ) | ||
| return b |
Comment on lines
195
to
208
| if options.start_qpos is not None: | ||
| start = options.start_qpos | ||
| if start.dim() == 1: | ||
| start = start.unsqueeze(0) | ||
| if qpos_list is not None: | ||
| qpos_list = torch.cat( | ||
| [options.start_qpos.unsqueeze(0), qpos_list], dim=0 | ||
| ) | ||
| qpos_list = torch.cat([start.unsqueeze(1), qpos_list], dim=1) | ||
| if xpos_list is not None: | ||
| start_xpos = self.robot.compute_fk( | ||
| qpos=options.start_qpos.unsqueeze(0), | ||
| name=options.control_part, | ||
| to_matrix=True, | ||
| qpos=start, name=options.control_part, to_matrix=True | ||
| ) | ||
| xpos_list = torch.cat([start_xpos, xpos_list], dim=0) | ||
| if start_xpos.dim() == 3: | ||
| start_xpos = start_xpos.unsqueeze(1) | ||
| xpos_list = torch.cat([start_xpos, xpos_list], dim=1) | ||
|
|
Comment on lines
215
to
220
| if not options.plan_opts: | ||
| # Directly return the interpolated trajectory if no further planning is needed | ||
| return PlanResult( | ||
| success=True, | ||
| positions=qpos_interpolated, | ||
| xpos_list=xpos_interpolated, | ||
| ) |
Comment on lines
+308
to
326
| motion_source = ( | ||
| getattr(cfg, "motion_source", "ik_interp") if cfg else "ik_interp" | ||
| ) | ||
| if motion_source == "motion_gen": | ||
| return self._plan_motion_gen( | ||
| target_states_list, | ||
| start_qpos, | ||
| n_waypoints, | ||
| control_part=control_part, | ||
| arm_dof=arm_dof, | ||
| cfg=cfg, | ||
| ) | ||
| return self._plan_ik_interp( | ||
| target_states_list, | ||
| start_qpos, | ||
| n_waypoints, | ||
| control_part=control_part, | ||
| is_interpolate=True, | ||
| is_linear=False, | ||
| interpolate_position_step=0.001, | ||
| plan_opts=ToppraPlanOptions(sample_interval=sample_interval), | ||
| arm_dof=arm_dof, | ||
| ) |
Comment on lines
+456
to
+473
| def _build_plan_opts(self, cfg: "ActionCfg | None", n_waypoints: int): | ||
| """Build planner options from action configuration.""" | ||
| planner_type = getattr(cfg, "planner_type", None) | ||
| if planner_type in (None, "toppra"): | ||
| constraints = {} | ||
| vl = getattr(cfg, "velocity_limit", None) | ||
| al = getattr(cfg, "acceleration_limit", None) | ||
| constraints["velocity"] = vl if vl is not None else 0.2 | ||
| constraints["acceleration"] = al if al is not None else 0.5 | ||
| return ToppraPlanOptions( | ||
| sample_method=TrajectorySampleMethod.QUANTITY, | ||
| sample_interval=n_waypoints, | ||
| constraints=constraints, | ||
| ) | ||
| # neural: planner reads its own cfg; pass minimal options | ||
| from embodichain.lab.sim.planners.neural_planner import NeuralPlanOptions | ||
|
|
||
| return NeuralPlanOptions() |
…com/DexForce/EmbodiChain into worktree-parallel-motion-generation
Comment on lines
+120
to
+124
| if target_states is not None and hasattr(self, "robot"): | ||
| robot_num = getattr(self.robot, "num_instances", None) | ||
| _check_batch_consistency( | ||
| target_states, expected_b=robot_num, robot_num_instances=robot_num | ||
| ) |
Comment on lines
+45
to
+51
| def _infer_batch_size(target_states: list[PlanState]) -> int | None: | ||
| """Return the leading batch dim B of the first tensor found in target_states, or None if none.""" | ||
| for s in target_states: | ||
| for t in (s.qpos, s.xpos, s.qvel, s.qacc): | ||
| if isinstance(t, torch.Tensor) and t.dim() >= 1: | ||
| return int(t.shape[0]) | ||
| return None |
Comment on lines
+119
to
+124
| target_states = kwargs.get("target_states", args[0] if args else None) | ||
| if target_states is not None and hasattr(self, "robot"): | ||
| robot_num = getattr(self.robot, "num_instances", None) | ||
| _check_batch_consistency( | ||
| target_states, expected_b=robot_num, robot_num_instances=robot_num | ||
| ) |
Comment on lines
115
to
+120
| def move_robot_along_trajectory( | ||
| robot: Robot, arm_name: str, qpos_list: list[torch.Tensor], delay: float = 0.1 | ||
| ): | ||
| """ | ||
| Set the robot joint positions sequentially along the given joint trajectory. | ||
| sim: SimulationManager, | ||
| robot: Robot, | ||
| arm_name: str, | ||
| qpos_trajectory: torch.Tensor | Sequence[torch.Tensor], | ||
| ) -> None: |
Comment on lines
+151
to
+153
| for qpos_step in qpos_trajectory.transpose(0, 1): | ||
| robot.set_qpos(qpos=qpos_step, joint_ids=joint_ids) | ||
| sim.update(step=4) |
Comment on lines
+285
to
+290
| move_robot_along_trajectory( | ||
| sim=sim, | ||
| robot=robot, | ||
| arm_name=arm_name, | ||
| qpos_trajectory=joint_plan.positions, | ||
| ) |
Comment on lines
+302
to
307
| move_robot_along_trajectory( | ||
| sim=sim, | ||
| robot=robot, | ||
| arm_name=arm_name, | ||
| qpos_trajectory=cartesian_plan.positions, | ||
| ) |
Comment on lines
+45
to
+51
| def _infer_batch_size(target_states: list[PlanState]) -> int | None: | ||
| """Return the leading batch dim B of the first tensor found in target_states, or None if none.""" | ||
| for s in target_states: | ||
| for t in (s.qpos, s.xpos, s.qvel, s.qacc): | ||
| if isinstance(t, torch.Tensor) and t.dim() >= 1: | ||
| return int(t.shape[0]) | ||
| return None |
Comment on lines
+129
to
133
| arm_name: Name of the robot arm to control. | ||
| qpos_trajectory: Joint positions shaped ``(B, N, DOF)``, ``(N, DOF)``, | ||
| or a sequence of waypoint tensors. | ||
| delay: Time delay between each step in seconds. | ||
| """ |
Comment on lines
+80
to
+83
| if sample_method == TrajectorySampleMethod.TIME and sample_interval <= 0: | ||
| return _empty_failure(dofs) | ||
| if sample_method == TrajectorySampleMethod.QUANTITY and sample_interval < 2: | ||
| return _empty_failure(dofs) |
Comment on lines
215
to
219
| if not options.plan_opts: | ||
| # Directly return the interpolated trajectory if no further planning is needed | ||
| return PlanResult( | ||
| success=True, | ||
| positions=qpos_interpolated, | ||
| xpos_list=xpos_interpolated, |
Comment on lines
447
to
+454
| dt = torch.full( | ||
| (positions_t.shape[0],), | ||
| float(self.cfg.dt), | ||
| dtype=torch.float32, | ||
| device=self.device, | ||
| ) | ||
| dt = dt.unsqueeze(0).expand(b, -1) | ||
| positions_t = positions_t.permute(1, 0, 2) |
Comment on lines
+119
to
+124
| target_states = kwargs.get("target_states", args[0] if args else None) | ||
| if target_states is not None and hasattr(self, "robot"): | ||
| robot_num = getattr(self.robot, "num_instances", None) | ||
| _check_batch_consistency( | ||
| target_states, expected_b=robot_num, robot_num_instances=robot_num | ||
| ) |
…com/DexForce/EmbodiChain into worktree-parallel-motion-generation
…com/DexForce/EmbodiChain into worktree-parallel-motion-generation
Comment on lines
+119
to
+124
| target_states = kwargs.get("target_states", args[0] if args else None) | ||
| if target_states is not None and hasattr(self, "robot"): | ||
| robot_num = getattr(self.robot, "num_instances", None) | ||
| _check_batch_consistency( | ||
| target_states, expected_b=robot_num, robot_num_instances=robot_num | ||
| ) |
Comment on lines
215
to
219
| if not options.plan_opts: | ||
| # Directly return the interpolated trajectory if no further planning is needed | ||
| return PlanResult( | ||
| success=True, | ||
| positions=qpos_interpolated, | ||
| xpos_list=xpos_interpolated, |
Comment on lines
+129
to
133
| arm_name: Name of the robot arm to control. | ||
| qpos_trajectory: Joint positions shaped ``(B, N, DOF)``, ``(N, DOF)``, | ||
| or a sequence of waypoint tensors. | ||
| delay: Time delay between each step in seconds. | ||
| """ |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
This PR adds env-batched
(B=num_envs, N, DOF)parallel motion generation across the planner layer, theMotionGeneratorfacade, and the atomic-action layer — making motion generation and atomic actions fully parallel alongside the already-batched IK solvers.Both motion sources (planner-based
MotionGeneratorand IK+interpolation) are available to the atomic-action layer via a per-actionmotion_sourcecfg.Design spec:
docs/superpowers/specs/2026-07-02-parallel-motion-generation-design.mdImplementation plan:
docs/superpowers/plans/2026-07-02-parallel-motion-generation.mdKey changes
planners/utils.py):PlanState/PlanResultgain a leadingBdim;PlanState.from_qpos/from_xpos/singlectors;PlanResult.is_all_success.planners/base_planner.py):_infer_batch_size+_check_batch_consistencywired into@validate_plan_options.planners/toppra_planner.py): pure-numpy module-level_toppra_solve_one_envworker; lazy fork-contextProcessPoolExecutorfan-out acrossBenvs (inline fallback forB=1);_assemble_batched_resultwith tail-padding forTIME-sampled trajectories of varying duration;BrokenProcessPoolhandling (close + rebuild).max_workers/mp_contextcfg.planners/neural_planner.py): natively batchedplan()/_parse_waypoints/_build_obs/_is_active_reached; per-env early-convergence holds (converged envs' action masked to zero so qpos does not drift);num_instances>1guard removed.planners/motion_generator.py): batchedgenerate()(start_qpos:(B,DOF)); batchedinterpolate_trajectory(joint-space + Cartesian); newinterpolate_xpos_batchedhelper.atomic_actions/):ActionCfg.motion_source/planner_type;ActionResult.success→(B,)tensor withsuccess_all/__bool__deprecation shim;TrajectoryBuilder.plan_arm_trajmotion-source strategy (ik_interpdefault,motion_genopt-in) with_to_batched_plan_statesadapter and_build_plan_opts;AtomicActionEngine.runper-env failure propagation (failed envs hold their last successful qpos); all 6 actions migrated to per-env tensor success.atom_action_utils.plan_trajectory, 6 atomic-action tutorials, the neural-planner benchmark, and 3 tableware task files migrated to the batched contract.Out of scope (deferred)
action_bank/configurable_action.pytask-graph layer (spec §8) — itsplan_trajectorypath is not yet migrated; the 4 action_bank unit tests pass (the path is untested).MoveJointswithmotion_source="motion_gen"; per-env divergentmove_type; fullaction_bank/integration with batched atomic actions.Type of change
Note on backward compatibility: The
plan()/generate()API was changed in-place to be natively batched (per the design decision D5). Existing single-env callers were migrated toPlanState.single(...)/[0]-indexing.ActionResult.__bool__emits aDeprecationWarningto ease the transition.Test Plan
tests/sim/planners/— 36 tests (batched PlanState/PlanResult, validation, TOPPRA worker + batched plan + pool-path + tail-pad + numerical-regression anchor, neural batched parse/rollout/early-convergence-hold, MotionGenerator batched generate + interpolate)tests/sim/atomic_actions/— 103 tests (ActionResult tensor, TrajectoryBuilder motion_source, engine per-env failure propagation, all 6 actions migrated, reach-equivalence e2e forik_interpandmotion_gen+TOPPRA within 2 cm in real sim)tests/benchmark/— 10 tests (neural-planner benchmark migrated)tests/gym/— 142 tests (no regressions; action_bank path untested-but-passing)black .run; all touched files formatted.Total: 149+ tests passing in the core suites.
🤖 Generated with Claude Code