diff --git a/agent_context/topics/motion-planning/planner-details.md b/agent_context/topics/motion-planning/planner-details.md index b1ea0b513..ab7889b70 100644 --- a/agent_context/topics/motion-planning/planner-details.md +++ b/agent_context/topics/motion-planning/planner-details.md @@ -104,6 +104,14 @@ the graph. Validate complete observation-layout metadata/fingerprint before rollout; model dimensions and waypoint capacity belong to the export contract. Inspect the `policy-deploy` package extra for runtime dependencies. +`NeuralPlannerCfg.dt` is bookkeeping, not solved timing: the rollout integrates +a joint delta per step, so a saturated step implies `action_scale / dt` rad/s +and overshoots a typical arm's velocity limit several times over. Setting +`constraints` swaps that nominal timing for a TOPPRA parameterization through +the shared entry point below; it resamples the path, so positions and poses +change with the timing, and it cannot reduce the path's own jerk. Default +remains `None`, which preserves the nominal-timing behavior. + Runtime/training frame differences must be expressed through explicit policy frame and TCP transforms. Target/FK quaternions follow the shared [simulation convention](../simulation-system/simulation-system.md); do not @@ -123,6 +131,14 @@ semantics are owned by [collision worlds](collision-worlds.md). ## Retiming and playback +`toppra_planner.py::retime_joint_paths()` is the shared entry point for +re-parameterizing an already-planned joint path under velocity and acceleration +limits, used by planners that emit geometry without executable timing. It is +the planner's own time parameterization with the same backend selection, so +CUDA paths and gradient requests run on Warp; it resamples along the fitted +spline, so recompute derived poses from its output rather than reusing the +input samples. + `compute/trajectory/timing.py` owns time-domain differentiation/resampling and `retime_to_control_grid()`. Zero-time position changes are invalid; repeated-time unchanged samples may represent padding/junctions. Time resampling preserves diff --git a/docs/source/api_reference/public_api.rst b/docs/source/api_reference/public_api.rst index d4b037867..ea1e1c533 100644 --- a/docs/source/api_reference/public_api.rst +++ b/docs/source/api_reference/public_api.rst @@ -1176,11 +1176,17 @@ embodichain.lab.sim.motion.planners.toppra_planner .. currentmodule:: embodichain.lab.sim.motion.planners.toppra_planner +Time-optimal joint-path parameterization. The planner solves timing for +waypoint input; :func:`retime_joint_paths` applies the same parameterization to +paths another planner already produced, such as a closed-loop policy rollout +whose samples carry only nominal timing. + .. autosummary:: ToppraPlanner ToppraPlannerCfg ToppraPlanOptions + retime_joint_paths embodichain.lab.sim.motion.planners.se3 --------------------------------------- diff --git a/embodichain/lab/sim/motion/planners/neural_planner.py b/embodichain/lab/sim/motion/planners/neural_planner.py index 3caf83b98..bf9a91dd0 100644 --- a/embodichain/lab/sim/motion/planners/neural_planner.py +++ b/embodichain/lab/sim/motion/planners/neural_planner.py @@ -30,7 +30,12 @@ _infer_batch_size, validate_plan_options, ) -from embodichain.lab.sim.motion.planners.utils import MoveType, PlanResult, PlanState +from embodichain.lab.sim.motion.planners.utils import ( + MoveType, + PlanResult, + PlanState, + TrajectorySampleMethod, +) from embodichain.utils import configclass, logger from embodichain.utils.math import quat_error_magnitude, quat_from_matrix @@ -316,7 +321,30 @@ class NeuralPlannerCfg(BasePlannerCfg): """ dt: float = 0.01 - """Nominal timestep reported in PlanResult.""" + """Output sampling period in seconds. + + Without :attr:`constraints` this is bookkeeping rather than executable + timing. + + The rollout integrates a joint delta per step and never solves a duration, + so this value is bookkeeping rather than executable timing: at the default + ``action_scale`` a saturated step implies ``action_scale / dt`` rad/s, + which exceeds a typical arm's joint velocity limit several times over. Set + :attr:`constraints` to replace it with a solved time parameterization, in + which case this becomes the period that trajectory is sampled on. + """ + + constraints: dict[str, float | list[float]] | None = None + """Optional ``velocity`` and ``acceleration`` limits for output retiming. + + ``None`` keeps the rollout's nominal timing and current behavior. When set, + the closed-loop samples are treated as a geometric path and re-parameterized + under these limits, so the returned timing and derivatives are executable. + Keys and value shapes match :class:`ToppraPlanOptions`: a scalar or one + value per controlled joint. Retiming resamples the path, so it changes the + reported positions and poses as well as the timing, and it cannot reduce + the path's own jerk -- only the duration over which it is traversed. + """ @configclass @@ -583,10 +611,26 @@ def plan( dt[:, 0] = 0.0 positions_t = positions_t.permute(1, 0, 2) xpos_t = xpos_t.permute(1, 0, 2, 3) + success = active_idx >= episode_k + if self.cfg.constraints is not None: + return self._retimed_result( + success, + positions_t, + control_part, + waypoints=( + waypoints_pos, + waypoints_quat, + waypoints_joint, + pos_mask, + rot_mask, + joint_mask, + ), + episode_k=episode_k, + qpos_limits=(lower, upper), + ) velocities_t, accelerations_t = self._compute_vel_acc_via_finite_diff( positions_t, dt ) - success = active_idx >= episode_k return PlanResult( success=success, positions=positions_t, @@ -596,6 +640,141 @@ def plan( dt=dt, ) + def _retimed_result( + self, + success: torch.Tensor, + positions: torch.Tensor, + control_part: str, + *, + waypoints: tuple[torch.Tensor, ...], + episode_k: int, + qpos_limits: tuple[torch.Tensor, torch.Tensor], + ) -> PlanResult: + """Re-parameterize the rollout path under the configured limits. + + Retiming fits a spline through the rollout samples and resamples it, so + the returned trajectory is not the one the rollout verified. Success is + therefore re-derived from the samples actually returned rather than + carried over: the resampled grid can step past a waypoint the rollout + stopped on, and a spline through samples clamped at a joint limit can + overshoot that limit between them. A row is reported as successful only + when its returned trajectory parameterizes, moves in nonzero time, + stays inside the joint limits, and still reaches every waypoint. + + Args: + success: Per-env rollout convergence of shape ``(B,)``. + positions: Rollout joint samples of shape ``(B, N, DOF)``. + control_part: Robot control part used for forward kinematics. + waypoints: Parsed waypoint targets and masks from + :meth:`_parse_waypoints`. + episode_k: Number of waypoints this episode must reach. + qpos_limits: Lower and upper joint-position limits. + + Returns: + PlanResult with solved timing, derivatives and recomputed poses. + """ + from .toppra_planner import retime_joint_paths + + # Sample the solved trajectory on the configured control period rather + # than on the rollout's step count. The rollout emits however many + # steps it happened to take, which is far too coarse once the duration + # stretches: a waypoint the rollout stopped on then falls between two + # output samples and is reported as missed. + retimed = retime_joint_paths( + positions, + constraints=dict(self.cfg.constraints), + sample_method=TrajectorySampleMethod.TIME, + sample_interval=float(self.cfg.dt), + device=self.device, + ) + retimed_positions = retimed.positions + # One batched FK over every sample; the waypoint re-check reuses it. + poses = self.robot.compute_batch_fk( + qpos=retimed_positions, name=control_part, to_matrix=True + ) + success = success & retimed.success.to(success.device) + success = success & self._retimed_is_executable( + retimed_positions, retimed.dt, qpos_limits + ) + success = success & self._retimed_reaches_waypoints( + retimed_positions, poses, waypoints, episode_k + ) + return PlanResult( + success=success, + positions=retimed_positions, + velocities=retimed.velocities, + accelerations=retimed.accelerations, + xpos_list=poses, + dt=retimed.dt, + ) + + def _retimed_is_executable( + self, + positions: torch.Tensor, + dt: torch.Tensor, + qpos_limits: tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + """Return whether each retimed row has real timing and legal positions. + + The shared retiming kernel returns a zero-duration result for a path + whose endpoints nearly coincide, which would claim a move with no time + to execute it. Spline interpolation through samples the rollout clamped + at a joint limit can also overshoot that limit between them, since the + fit sees only the sampled values. + """ + lower, upper = qpos_limits + arm = positions[..., : self._action_dim] + moves = (arm - arm[:, :1]).abs().amax(dim=(1, 2)) > 1.0e-9 + timed = dt.sum(dim=-1) > 0.0 + within = ((arm >= lower) & (arm <= upper)).all(dim=-1).all(dim=-1) + return within & (timed | ~moves) + + def _retimed_reaches_waypoints( + self, + positions: torch.Tensor, + poses: torch.Tensor, + waypoints: tuple[torch.Tensor, ...], + episode_k: int, + ) -> torch.Tensor: + """Return whether each retimed row still reaches every waypoint in order. + + Applies the rollout's own arrival test to the resampled grid, so a + trajectory that passes a waypoint only between output samples is not + reported as reaching it. + """ + ( + waypoints_pos, + waypoints_quat, + waypoints_joint, + pos_mask, + rot_mask, + joint_mask, + ) = waypoints + batch, samples = positions.shape[:2] + policy_poses = self._policy_pose_xyzw(poses.flatten(0, 1)).view( + batch, samples, -1 + ) + active_idx = torch.zeros(batch, dtype=torch.long, device=self.device) + for index in range(samples): + qpos = positions[:, index] + reached = self._is_active_reached( + qpos[:, : self._action_dim], + policy_poses[:, index], + waypoints_pos, + waypoints_quat, + waypoints_joint, + pos_mask, + rot_mask, + joint_mask, + active_idx, + ) + active_idx = torch.where( + reached & (active_idx < episode_k), active_idx + 1, active_idx + ) + if bool((active_idx >= episode_k).all()): + break + return active_idx >= episode_k + def _parse_waypoints(self, target_states: list[PlanState]) -> tuple[ torch.Tensor, torch.Tensor, @@ -702,7 +881,11 @@ def _fk_matrix(self, qpos: torch.Tensor, control_part: str) -> torch.Tensor: def _fk_pose_xyzw(self, qpos: torch.Tensor, control_part: str) -> torch.Tensor: """Return the policy-frame FK pose as ``xyz + xyzw``.""" - fk = self._to_policy_frame(self._fk_matrix(qpos, control_part)) + return self._policy_pose_xyzw(self._fk_matrix(qpos, control_part)) + + def _policy_pose_xyzw(self, fk_matrix: torch.Tensor) -> torch.Tensor: + """Map ``(M, 4, 4)`` runtime TCP poses to policy-frame ``xyz + xyzw``.""" + fk = self._to_policy_frame(fk_matrix) pos = fk[:, :3, 3] # ``quat_from_matrix`` is an EmbodiChain ``xyzw`` producer; converting # it again would turn a valid pose into a different rotation. diff --git a/embodichain/lab/sim/motion/planners/toppra_planner.py b/embodichain/lab/sim/motion/planners/toppra_planner.py index 1cf21e21c..49bad508b 100644 --- a/embodichain/lab/sim/motion/planners/toppra_planner.py +++ b/embodichain/lab/sim/motion/planners/toppra_planner.py @@ -19,6 +19,7 @@ from copy import deepcopy import os +from collections.abc import Mapping, Sequence from typing import Literal import torch @@ -174,7 +175,171 @@ def _worker_init() -> None: _set_parent_death_signal() -__all__ = ["ToppraPlanner", "ToppraPlannerCfg", "ToppraPlanOptions"] +def _assemble_batched_results( + results: list[dict], + dofs: int, + device: torch.device, +) -> PlanResult: + """Stack per-env solves into one batched :class:`PlanResult`. + + Env trajectories may differ in length. Shorter rows are tail-padded by + repeating their final waypoint with zero velocity and acceleration, so + every output tensor shares the same ``(B, N, DOF)`` or ``(B, N)`` shape. + + Args: + results: Per-env dicts as returned by :func:`_toppra_solve_one_env`. + dofs: Per-env degrees of freedom. + device: Device for the assembled tensors. + + Returns: + PlanResult with env-batched positions, derivatives, timing and success. + """ + b = len(results) + max_n = max(r["n"] for r in results) + positions = np.zeros((b, max_n, dofs), dtype=np.float32) + velocities = np.zeros((b, max_n, dofs), dtype=np.float32) + accelerations = np.zeros((b, max_n, dofs), dtype=np.float32) + dt = np.zeros((b, max_n), dtype=np.float32) + success = np.zeros((b,), dtype=bool) + for i, r in enumerate(results): + n = r["n"] + positions[i, :n] = r["positions"] + velocities[i, :n] = r["velocities"] + accelerations[i, :n] = r["accelerations"] + dt[i, :n] = r["dt"] + success[i] = r["success"] + # tail-pad: repeat final waypoint for held-pose rows + if n < max_n: + positions[i, n:] = r["positions"][-1] + velocities[i, n:] = 0.0 + accelerations[i, n:] = 0.0 + return PlanResult( + success=torch.as_tensor(success, device=device), + positions=torch.as_tensor(positions, device=device), + velocities=torch.as_tensor(velocities, device=device), + accelerations=torch.as_tensor(accelerations, device=device), + dt=torch.as_tensor(dt, device=device), + ) + + +def retime_joint_paths( + positions: torch.Tensor, + *, + constraints: Mapping[str, float | Sequence[float] | torch.Tensor], + sample_method: TrajectorySampleMethod = TrajectorySampleMethod.QUANTITY, + sample_interval: float | int | None = None, + grid_size: int = 100, + device: torch.device | None = None, +) -> PlanResult: + """Re-parameterize already-planned joint paths under dynamic limits. + + This is the time parameterization of :class:`ToppraPlanner`, exposed for + planners that produce geometry without executable timing. A closed-loop + policy rollout, for example, emits one joint sample per control step on a + nominal ``dt`` that no dynamic limit informed; feeding those samples here + replaces that timing with a solved one that respects the limits. + + Backend selection follows :class:`ToppraPlanner`'s ``auto`` mode: Warp on + CUDA inputs or when a gradient is required, NumPy otherwise. Both enforce + the limits over every grid interval rather than only at grid points. + + The path geometry is resampled along the fitted spline, so poses derived + from the input samples must be recomputed from the returned positions. + + Args: + positions: Planned joint paths of shape ``(B, N, DOF)``. + constraints: ``velocity`` and ``acceleration`` limits, each a scalar, + one value per joint, or signed ``(DOF, 2)`` bounds, matching + :class:`ToppraPlanOptions`. + sample_method: Fixed output quantity or approximately fixed time step. + sample_interval: Output count for ``QUANTITY`` or seconds for ``TIME``. + ``None`` keeps the input sample count and is only meaningful for + ``QUANTITY``, which leaves array shapes unchanged while the timing + behind them changes. + grid_size: Minimum parameterization grid size, as for + :class:`ToppraPlannerCfg`. + device: Device for the assembled tensors; defaults to the input's. + + Returns: + PlanResult whose ``success`` is per-env and false where the path could + not be parameterized. + + Raises: + ValueError: If ``positions`` is not a ``(B, N, DOF)`` tensor, the + constraints omit ``velocity`` or ``acceleration``, or ``TIME`` + sampling is requested without an explicit interval in seconds. + """ + if positions.dim() != 3: + raise ValueError( + f"positions must have shape (B, N, DOF), got {tuple(positions.shape)}." + ) + missing = sorted({"velocity", "acceleration"}.difference(constraints)) + if missing: + raise ValueError(f"constraints is missing required keys: {missing}.") + if sample_interval is None and sample_method is TrajectorySampleMethod.TIME: + # The default is a sample count; reading it as seconds would silently + # collapse a dense path to two samples whenever its duration is shorter. + raise ValueError( + "sample_interval is required for TIME sampling; it is a duration in " + "seconds, not the sample count the default supplies." + ) + resolved_device = positions.device if device is None else device + batch, count, dofs = positions.shape + resolved_interval = count if sample_interval is None else sample_interval + velocity, acceleration = constraints["velocity"], constraints["acceleration"] + needs_grad = torch.is_grad_enabled() and ( + positions.requires_grad + or any( + isinstance(limit, torch.Tensor) and limit.requires_grad + for limit in (velocity, acceleration) + ) + ) + if positions.is_cuda or needs_grad: + from embodichain.compute.trajectory._toppra_warp import _retime_toppra_warp + + quantity = sample_method == TrajectorySampleMethod.QUANTITY + waypoints = positions if positions.is_floating_point() else positions.double() + try: + result = _retime_toppra_warp( + waypoints, + velocity, + acceleration, + sample_count=int(resolved_interval) if quantity else None, + sample_dt=None if quantity else float(resolved_interval), + grid_size=grid_size, + ) + except (ValueError, FloatingPointError, OverflowError): + return _assemble_batched_results( + [_empty_failure(dofs) for _ in range(batch)], dofs, resolved_device + ) + return PlanResult( + **{key: value.to(resolved_device) for key, value in result.items()} + ) + + samples = positions.detach().cpu().numpy().astype(np.float64) + to_numpy = lambda limit: ( + limit.detach().cpu().numpy() if isinstance(limit, torch.Tensor) else limit + ) + results = [ + _toppra_solve_one_env( + samples[index], + to_numpy(velocity), + to_numpy(acceleration), + sample_method, + resolved_interval, + grid_size, + ) + for index in range(batch) + ] + return _assemble_batched_results(results, dofs, resolved_device) + + +__all__ = [ + "ToppraPlanner", + "ToppraPlannerCfg", + "ToppraPlanOptions", + "retime_joint_paths", +] @configclass @@ -535,47 +700,5 @@ def plan( return self._assemble_batched_result(results, dofs) def _assemble_batched_result(self, results: list[dict], dofs: int) -> PlanResult: - """Stack per-env TOPPRA results into a batched :class:`PlanResult`. - - Each entry of ``results`` is the dict returned by - :func:`_toppra_solve_one_env`. Env trajectories may have different - lengths (``n``); this method pads shorter trajectories out to the - longest by repeating their final waypoint (held pose) with zero - velocity and acceleration, so every output tensor shares the same - ``(B, N, DOF)`` / ``(B, N)`` shape. - - Args: - results: list of per-env result dicts (length ``B``). - dofs: per-env degrees of freedom. - - Returns: - PlanResult with env-batched tensors (``success`` ``(B,)``, - ``positions``/``velocities``/``accelerations`` ``(B, N, DOF)``, - ``dt`` ``(B, N)``, ``duration`` ``(B,)``). - """ - b = len(results) - max_n = max(r["n"] for r in results) - positions = np.zeros((b, max_n, dofs), dtype=np.float32) - velocities = np.zeros((b, max_n, dofs), dtype=np.float32) - accelerations = np.zeros((b, max_n, dofs), dtype=np.float32) - dt = np.zeros((b, max_n), dtype=np.float32) - success = np.zeros((b,), dtype=bool) - for i, r in enumerate(results): - n = r["n"] - positions[i, :n] = r["positions"] - velocities[i, :n] = r["velocities"] - accelerations[i, :n] = r["accelerations"] - dt[i, :n] = r["dt"] - success[i] = r["success"] - # tail-pad: repeat final waypoint for held-pose rows - if n < max_n: - positions[i, n:] = r["positions"][-1] - velocities[i, n:] = 0.0 - accelerations[i, n:] = 0.0 - return PlanResult( - success=torch.as_tensor(success, device=self.device), - positions=torch.as_tensor(positions, device=self.device), - velocities=torch.as_tensor(velocities, device=self.device), - accelerations=torch.as_tensor(accelerations, device=self.device), - dt=torch.as_tensor(dt, device=self.device), - ) + """Stack per-env TOPPRA results into a batched :class:`PlanResult`.""" + return _assemble_batched_results(results, dofs, self.device) diff --git a/tests/sim/motion/planners/test_neural_batched.py b/tests/sim/motion/planners/test_neural_batched.py index e2ebc039c..df263aeee 100644 --- a/tests/sim/motion/planners/test_neural_batched.py +++ b/tests/sim/motion/planners/test_neural_batched.py @@ -87,6 +87,7 @@ def test_plan_returns_batched_success(self, monkeypatch): "dt": 0.01, "control_part": "arm", "num_arm_joints": 7, + "constraints": None, }, )() @@ -158,6 +159,7 @@ def test_converged_env_holds_qpos(self): "dt": 0.01, "control_part": "arm", "num_arm_joints": 7, + "constraints": None, }, )() diff --git a/tests/sim/motion/planners/test_neural_planner.py b/tests/sim/motion/planners/test_neural_planner.py index 73f91caa7..82db1f152 100644 --- a/tests/sim/motion/planners/test_neural_planner.py +++ b/tests/sim/motion/planners/test_neural_planner.py @@ -31,6 +31,8 @@ ) from embodichain.lab.sim.motion.planners import neural_planner as neural_planner_module from embodichain.lab.sim.motion.planners.neural_planner import NeuralPlanOptions +from embodichain.lab.sim.motion.planners.toppra_planner import retime_joint_paths +from embodichain.lab.sim.motion.planners.utils import TrajectorySampleMethod from embodichain.lab.sim.sim_manager import SimulationManager NUM_ARM_JOINTS = 7 @@ -116,6 +118,13 @@ def compute_fk( return torch.tensor([[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]]).repeat(batch, 1) +def _fake_batch_fk(self, qpos, name=None, env_ids=None, to_matrix=False): + return torch.eye(4).repeat(qpos.shape[0], qpos.shape[1], 1, 1) + + +FakeRobot.compute_batch_fk = _fake_batch_fk + + class FakeSimulationManager: def __init__(self): self.robot = FakeRobot() @@ -402,6 +411,134 @@ def test_neural_planner_returns_velocities_and_accelerations(tmp_path, monkeypat assert torch.isfinite(result.accelerations).all() +class SaturatingOnnxPolicy(FakeOnnxPolicy): + """Policy that always saturates, the worst case for implied velocity.""" + + def __call__(self, obs: torch.Tensor) -> torch.Tensor: + self.last_obs = obs.clone() + return torch.ones(obs.shape[0], NUM_ARM_JOINTS, device=obs.device) + + +def _rollout(tmp_path, monkeypatch, *, policy=None, **cfg_overrides): + """Roll the neural planner out toward one identity pose target.""" + model_path = _create_fake_onnx_model(tmp_path) + fake_sim = FakeSimulationManager() + monkeypatch.setattr( + SimulationManager, + "get_instance", + classmethod(lambda cls, instance_id=0: fake_sim), + ) + if policy is not None: + monkeypatch.setattr(neural_planner_module, "_OnnxPolicy", policy) + motion_generator = MotionGenerator( + cfg=MotionGenCfg( + planner_cfg=NeuralPlannerCfg( + robot_uid="fake_robot", + onnx_model_path=model_path, + control_part="main_arm", + **cfg_overrides, + ) + ) + ) + return motion_generator.generate( + target_states=[ + PlanState.single(move_type=MoveType.EEF_MOVE, xpos=torch.eye(4)) + ], + options=MotionGenOptions( + plan_opts=NeuralPlanOptions( + control_part="main_arm", + start_qpos=torch.zeros(NUM_ARM_JOINTS), + ), + ), + ) + + +def test_neural_planner_keeps_nominal_timing_without_constraints(tmp_path, monkeypatch): + result = _rollout(tmp_path, monkeypatch, policy=SaturatingOnnxPolicy) + + # The default path is unchanged: one nominal interval per rollout step. + intervals = result.dt[0, 1:] + assert torch.allclose(intervals, torch.full_like(intervals, 0.01)) + # 0.2 rad per 0.01 s is 20 rad/s, which no arm can execute. + assert float(result.velocities.abs().max()) == pytest.approx(20.0, rel=1e-3) + + +def test_neural_planner_constraints_replace_timing_and_bound_dynamics( + tmp_path, monkeypatch +): + nominal = _rollout(tmp_path, monkeypatch, policy=SaturatingOnnxPolicy) + retimed = _rollout( + tmp_path, + monkeypatch, + policy=SaturatingOnnxPolicy, + constraints={"velocity": 1.0, "acceleration": 2.0}, + ) + + assert bool(retimed.success.all()) + assert float(retimed.velocities.abs().max()) <= 1.0 + 1e-6 + assert float(retimed.accelerations.abs().max()) <= 2.0 + 1e-6 + # Respecting the limits costs time; the nominal timing understated it. + assert float(retimed.dt.sum()) > float(nominal.dt.sum()) + # Retiming resamples the path, so poses must track the returned positions. + assert retimed.xpos_list.shape[:2] == retimed.positions.shape[:2] + + +def test_retime_joint_paths_rejects_malformed_input(): + positions = torch.zeros(1, 4, NUM_ARM_JOINTS) + limits = {"velocity": 1.0, "acceleration": 1.0} + with pytest.raises(ValueError, match="velocity"): + retime_joint_paths(positions, constraints={"acceleration": 1.0}) + with pytest.raises(ValueError, match=r"\(B, N, DOF\)"): + retime_joint_paths(positions[0], constraints=limits) + # The default interval is a sample count; reading it as seconds would + # silently collapse a dense path to two samples. + with pytest.raises(ValueError, match="TIME sampling"): + retime_joint_paths( + positions, + constraints=limits, + sample_method=TrajectorySampleMethod.TIME, + ) + + +def test_retiming_samples_on_the_control_period_not_the_rollout_step_count( + tmp_path, monkeypatch +): + # The rollout emits however many steps it happened to take. Inheriting that + # count leaves a stretched trajectory sampled far too coarsely, so a + # waypoint the rollout stopped on falls between two output samples. + nominal = _rollout(tmp_path, monkeypatch, policy=SaturatingOnnxPolicy) + retimed = _rollout( + tmp_path, + monkeypatch, + policy=SaturatingOnnxPolicy, + constraints={"velocity": 1.0, "acceleration": 2.0}, + ) + duration = float(retimed.dt.sum()) + assert duration > float(nominal.dt.sum()) + # One sample per nominal dt over the solved duration, not one per step. + assert retimed.positions.shape[1] > nominal.positions.shape[1] + assert retimed.positions.shape[1] == pytest.approx(duration / 0.01, rel=0.05) + + +def test_retimed_success_reflects_the_trajectory_actually_returned( + tmp_path, monkeypatch +): + # A joint limit the rollout clamps against is also a limit the spline must + # respect, and the resampled grid must still reach the goal. Success is + # re-derived from the returned samples rather than inherited. + retimed = _rollout( + tmp_path, + monkeypatch, + policy=SaturatingOnnxPolicy, + constraints={"velocity": 1.0, "acceleration": 2.0}, + ) + arm = retimed.positions[..., :NUM_ARM_JOINTS] + # FakeRobot declares +-2.0 rad on every joint. + assert float(arm.max()) <= 2.0 + 1e-6 + assert float(arm.min()) >= -2.0 - 1e-6 + assert float(retimed.dt.sum()) > 0.0 + + def test_neural_planner_disables_grad_for_all_fk_calls(tmp_path, monkeypatch): model_path = _create_fake_onnx_model(tmp_path) fake_sim = FakeSimulationManager()