refactor(manipulation): extract execution runtime

dimos/control/tasks/trajectory_task/trajectory_task.py

@@ -23,6 +23,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+import math
 from typing import Any
 
 from dimos.control.task import (
@@ -127,7 +128,7 @@ def compute(self, state: CoordinatorState) -> JointCommandOutput | None:
         Returns:
             JointCommandOutput with positions, or None if not executing
         """
-        if self._trajectory is None:
+        if self._trajectory is None or not self._trajectory.joint_names:
             return None
 
         # Set start time on first compute() for consistent timing
@@ -143,8 +144,10 @@ def compute(self, state: CoordinatorState) -> JointCommandOutput | None:
             logger.info(f"Trajectory {self._name} completed after {t_elapsed:.3f}s")
             # Return final position to hold at goal
             q_ref, _ = self._trajectory.sample(self._trajectory.duration)
+            final_names = list(self._trajectory.joint_names)
+            self._clear_active_trajectory()
             return JointCommandOutput(
-                joint_names=self._joint_names_list,
+                joint_names=final_names,
                 positions=list(q_ref),
                 mode=ControlMode.SERVO_POSITION,
             )
@@ -153,7 +156,7 @@ def compute(self, state: CoordinatorState) -> JointCommandOutput | None:
         q_ref, _ = self._trajectory.sample(t_elapsed)
 
         return JointCommandOutput(
-            joint_names=self._joint_names_list,
+            joint_names=list(self._trajectory.joint_names),
             positions=list(q_ref),
             mode=ControlMode.SERVO_POSITION,
         )
@@ -169,6 +172,60 @@ def on_preempted(self, by_task: str, joints: frozenset[str]) -> None:
         # Abort if any of our joints were preempted
         if joints & self._joint_names:
             self._state = TrajectoryState.ABORTED
+            self._clear_active_trajectory()
+
+    def _clear_active_trajectory(self) -> None:
+        """Clear stored trajectory-specific execution state."""
+        self._trajectory = None
+        self._pending_start = False
+        self._start_time = 0.0
+
+    def _validate_trajectory(self, trajectory: JointTrajectory) -> bool:
+        """Validate a trajectory before execution."""
+        joint_names = list(trajectory.joint_names)
+        if not joint_names:
+            logger.warning("Trajectory for %s has empty joint names", self._name)
+            return False
+        if len(set(joint_names)) != len(joint_names):
+            logger.warning("Trajectory for %s has duplicate joint names", self._name)
+            return False
+        unknown = [name for name in joint_names if name not in self._joint_names]
+        if unknown:
+            logger.warning("Trajectory for %s has unknown joints: %s", self._name, unknown)
+            return False
+        if not trajectory.points:
+            logger.warning("Empty trajectory for %s", self._name)
+            return False
+        width = len(joint_names)
+        previous_time: float | None = None
+        for index, point in enumerate(trajectory.points):
+            if len(point.positions) != width or len(point.velocities) != width:
+                logger.warning("Trajectory point %d for %s has invalid width", index, self._name)
+                return False
+            if not all(math.isfinite(value) for value in point.positions):
+                logger.warning(
+                    "Trajectory point %d for %s has non-finite positions", index, self._name
+                )
+                return False
+            if not all(math.isfinite(value) for value in point.velocities):
+                logger.warning(
+                    "Trajectory point %d for %s has non-finite velocities", index, self._name
+                )
+                return False
+            if not math.isfinite(point.time_from_start):
+                logger.warning("Trajectory point %d for %s has non-finite time", index, self._name)
+                return False
+            if index == 0 and point.time_from_start != 0.0:
+                logger.warning("Trajectory for %s must start at t=0", self._name)
+                return False
+            if previous_time is not None and point.time_from_start <= previous_time:
+                logger.warning("Trajectory for %s has non-increasing timestamps", self._name)
+                return False
+            previous_time = point.time_from_start
+        if trajectory.duration <= 0.0:
+            logger.warning("Trajectory for %s has nonpositive duration", self._name)
+            return False
+        return True
 
     def execute(self, trajectory: JointTrajectory) -> bool:
         """Start executing a trajectory.
@@ -183,17 +240,17 @@ def execute(self, trajectory: JointTrajectory) -> bool:
             logger.warning(f"Cannot execute: {self._name} in FAULT state")
             return False
 
-        if trajectory is None or trajectory.duration <= 0:
+        if trajectory is None:
             logger.warning(f"Invalid trajectory for {self._name}")
             return False
 
-        if not trajectory.points:
-            logger.warning(f"Empty trajectory for {self._name}")
+        if not self._validate_trajectory(trajectory):
             return False
 
         # Preempt any active trajectory
         if self._state == TrajectoryState.EXECUTING:
             logger.info(f"Preempting active trajectory on {self._name}")
+            self._clear_active_trajectory()
 
         self._trajectory = trajectory
         self._pending_start = True  # Start time set on first compute()
@@ -214,6 +271,7 @@ def cancel(self) -> bool:
         if self._state != TrajectoryState.EXECUTING:
             return False
         self._state = TrajectoryState.ABORTED
+        self._clear_active_trajectory()
         logger.info(f"Trajectory {self._name} cancelled")
         return True
 
@@ -227,7 +285,7 @@ def reset(self) -> bool:
             logger.warning(f"Cannot reset {self._name} while executing")
             return False
         self._state = TrajectoryState.IDLE
-        self._trajectory = None
+        self._clear_active_trajectory()
         logger.info(f"Trajectory {self._name} reset to IDLE")
         return True
 

dimos/control/test_control.py

@@ -428,6 +428,141 @@ def test_execute_trajectory(self, trajectory_task, simple_trajectory):
         assert trajectory_task.is_active()
         assert trajectory_task.get_state() == TrajectoryState.EXECUTING
 
+    def test_execute_partial_subset_and_claims_full_configuration(self, trajectory_task):
+        trajectory = JointTrajectory(
+            joint_names=["arm/joint2", "arm/joint3"],
+            points=[
+                TrajectoryPoint(positions=[0.0, 0.0], velocities=[0.0, 0.0], time_from_start=0.0),
+                TrajectoryPoint(positions=[0.5, 1.0], velocities=[0.0, 0.0], time_from_start=1.0),
+            ],
+        )
+
+        assert trajectory_task.execute(trajectory) is True
+        assert trajectory_task.claim().joints == frozenset(
+            {"arm/joint1", "arm/joint2", "arm/joint3"}
+        )
+
+    @pytest.mark.parametrize(
+        "trajectory",
+        [
+            JointTrajectory(
+                joint_names=[],
+                points=[TrajectoryPoint(time_from_start=0.0, positions=[], velocities=[])],
+            ),
+            JointTrajectory(
+                joint_names=["arm/joint1", "arm/joint1"],
+                points=[
+                    TrajectoryPoint(
+                        time_from_start=0.0, positions=[0.0, 0.0], velocities=[0.0, 0.0]
+                    )
+                ],
+            ),
+            JointTrajectory(
+                joint_names=["arm/missing"],
+                points=[TrajectoryPoint(time_from_start=0.0, positions=[0.0], velocities=[0.0])],
+            ),
+            JointTrajectory(joint_names=["arm/joint1"], points=[]),
+            JointTrajectory(
+                joint_names=["arm/joint1"],
+                points=[TrajectoryPoint(time_from_start=0.0, positions=[], velocities=[0.0])],
+            ),
+            JointTrajectory(
+                joint_names=["arm/joint1"],
+                points=[TrajectoryPoint(time_from_start=0.0, positions=[0.0], velocities=[0.0])],
+            ),
+            JointTrajectory(
+                joint_names=["arm/joint1", "arm/joint2", "arm/joint3"],
+                points=[
+                    TrajectoryPoint(
+                        time_from_start=0.0,
+                        positions=[0.0, 0.0, 0.0],
+                        velocities=[0.0, 0.0, 0.0],
+                    )
+                ],
+            ),
+            JointTrajectory(
+                joint_names=["arm/joint1"],
+                points=[
+                    TrajectoryPoint(time_from_start=0.0, positions=[float("nan")], velocities=[0.0])
+                ],
+            ),
+            JointTrajectory(
+                joint_names=["arm/joint1"],
+                points=[TrajectoryPoint(time_from_start=0.1, positions=[0.0], velocities=[0.0])],
+            ),
+            JointTrajectory(
+                joint_names=["arm/joint1"],
+                points=[
+                    TrajectoryPoint(time_from_start=0.0, positions=[0.0], velocities=[0.0]),
+                    TrajectoryPoint(time_from_start=0.0, positions=[1.0], velocities=[0.0]),
+                ],
+            ),
+        ],
+    )
+    def test_invalid_partial_inputs_reject_before_state_changes(self, trajectory_task, trajectory):
+        assert trajectory_task.get_state() == TrajectoryState.IDLE
+        assert trajectory_task.execute(trajectory) is False
+        assert trajectory_task.get_state() == TrajectoryState.IDLE
+        assert (
+            trajectory_task.compute(CoordinatorState(joints=MagicMock(), t_now=0.0, dt=0.01))
+            is None
+        )
+
+    def test_compute_emits_active_subset_only_and_clears_on_completion(self, trajectory_task):
+        trajectory = JointTrajectory(
+            joint_names=["arm/joint2"],
+            points=[
+                TrajectoryPoint(positions=[0.0], velocities=[0.0], time_from_start=0.0),
+                TrajectoryPoint(positions=[1.0], velocities=[0.0], time_from_start=1.0),
+            ],
+        )
+        assert trajectory_task.execute(trajectory) is True
+        trajectory_task.compute(CoordinatorState(joints=MagicMock(), t_now=10.0, dt=0.01))
+        output = trajectory_task.compute(CoordinatorState(joints=MagicMock(), t_now=10.5, dt=0.01))
+        assert output is not None
+        assert output.joint_names == ["arm/joint2"]
+        assert output.positions == [pytest.approx(0.5)]
+
+        final = trajectory_task.compute(CoordinatorState(joints=MagicMock(), t_now=11.5, dt=0.01))
+        assert final is not None
+        assert final.joint_names == ["arm/joint2"]
+        assert trajectory_task.get_state() == TrajectoryState.COMPLETED
+        assert (
+            trajectory_task.compute(CoordinatorState(joints=MagicMock(), t_now=12.0, dt=0.01))
+            is None
+        )
+
+    def test_replacement_reset_and_cancel_clear_active_subset(self, trajectory_task):
+        first = JointTrajectory(
+            joint_names=["arm/joint1"],
+            points=[
+                TrajectoryPoint(positions=[0.0], velocities=[0.0], time_from_start=0.0),
+                TrajectoryPoint(positions=[1.0], velocities=[0.0], time_from_start=1.0),
+            ],
+        )
+        second = JointTrajectory(
+            joint_names=["arm/joint3"],
+            points=[
+                TrajectoryPoint(positions=[2.0], velocities=[0.0], time_from_start=0.0),
+                TrajectoryPoint(positions=[3.0], velocities=[0.0], time_from_start=1.0),
+            ],
+        )
+        assert trajectory_task.execute(first) is True
+        assert trajectory_task.execute(second) is True
+        trajectory_task.compute(CoordinatorState(joints=MagicMock(), t_now=1.0, dt=0.01))
+        output = trajectory_task.compute(CoordinatorState(joints=MagicMock(), t_now=1.5, dt=0.01))
+        assert output is not None
+        assert output.joint_names == ["arm/joint3"]
+        assert trajectory_task.cancel() is True
+        assert (
+            trajectory_task.compute(CoordinatorState(joints=MagicMock(), t_now=2.0, dt=0.01))
+            is None
+        )
+        assert trajectory_task.reset() is True
+        assert trajectory_task.claim().joints == frozenset(
+            {"arm/joint1", "arm/joint2", "arm/joint3"}
+        )
+
     def test_compute_during_trajectory(self, trajectory_task, simple_trajectory, coordinator_state):
         t_start = time.perf_counter()
         trajectory_task.execute(simple_trajectory)