diff --git a/dimos/control/tasks/trajectory_task/test_trajectory_task.py b/dimos/control/tasks/trajectory_task/test_trajectory_task.py new file mode 100644 index 0000000000..cbef3f88cf --- /dev/null +++ b/dimos/control/tasks/trajectory_task/test_trajectory_task.py @@ -0,0 +1,133 @@ +# Copyright 2026 Dimensional Inc. +# +# 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. + +import pytest + +from dimos.control.task import ( + CoordinatorState, + JointStateSnapshot, +) +from dimos.control.tasks.trajectory_task.trajectory_task import ( + JointTrajectoryTask, + JointTrajectoryTaskConfig, + TrajectoryExecutionStatus, +) +from dimos.msgs.trajectory_msgs.JointTrajectory import JointTrajectory +from dimos.msgs.trajectory_msgs.TrajectoryPoint import TrajectoryPoint + + +@pytest.mark.parametrize("single_point", [False, True]) +@pytest.mark.parametrize("preempted", [False, True]) +def test_jtt_does_not_pull_joint_back_after_another_task_moves_it(single_point, preempted): + task = JointTrajectoryTask(JointTrajectoryTaskConfig(joint_names=["joint"])) + state = CoordinatorState( + joints=JointStateSnapshot(joint_positions={"joint": 0.0}), t_now=0.1, dt=0.1 + ) + first_target = 1.0 if preempted else 0.1 + first = JointTrajectory( + joint_names=["joint"], points=[TrajectoryPoint(positions=[first_target])] + ) + assert task.execute(first, {}).status is TrajectoryExecutionStatus.ACCEPTED + output = task.compute(state) + assert output is not None + assert output.positions == pytest.approx([0.1]) + if preempted: + task.on_preempted("other_task", frozenset({"joint"})) + assert not task.is_active() + + # Another task moves the joint away from JTT's previous command. + state = CoordinatorState( + joints=JointStateSnapshot(joint_positions={"joint": -1.0}), t_now=1.0, dt=0.01 + ) + assert task.compute(state) is None + points = [TrajectoryPoint(positions=[-1.2])] + if not single_point: + points = [ + TrajectoryPoint(positions=[-1.0]), + TrajectoryPoint(positions=[-1.2], time_from_start=0.1), + ] + assert ( + task.execute( + JointTrajectory(joint_names=["joint"], points=points), state.joints.joint_positions + ).status + is TrajectoryExecutionStatus.ACCEPTED + ) + + # The new motion goes farther away from the old command. Every output + # must move toward the new target, bounded from the new starting pose. + previous = -1.0 + for tick in range(30): + state.t_now = 1.0 + tick * state.dt + output = task.compute(state) + if output is None: + break + commanded = output.positions[0] + assert -state.dt - 1e-9 <= commanded - previous <= 1e-9 + previous = commanded + assert previous == pytest.approx(-1.2) + assert not task.is_active() + assert task.compute(state) is None + + +def test_completed_joint_reanchors_while_other_joint_keeps_command_continuity(): + task = JointTrajectoryTask(JointTrajectoryTaskConfig(joint_names=["finished", "running"])) + measured = JointStateSnapshot(joint_positions={"finished": 0.0, "running": 0.0}) + initial = JointTrajectory( + joint_names=["finished", "running"], points=[TrajectoryPoint(positions=[0.1, 1.0])] + ) + assert task.execute(initial, {}).status is TrajectoryExecutionStatus.ACCEPTED + output = task.compute(CoordinatorState(joints=measured, t_now=0.1, dt=0.1)) + assert output is not None + assert output.positions == pytest.approx([0.1, 0.1]) + measured.joint_positions["finished"] = -1.0 + replacement = JointTrajectory( + joint_names=["finished"], + points=[ + TrajectoryPoint(positions=[-1.0]), + TrajectoryPoint(positions=[0.1], time_from_start=2.0), + ], + ) + assert ( + task.execute(replacement, measured.joint_positions).status + is TrajectoryExecutionStatus.ACCEPTED + ) + output = task.compute(CoordinatorState(joints=measured, t_now=0.2, dt=0.1)) + assert output is not None + assert output.positions == pytest.approx([-1.0, 0.2]) + + +@pytest.mark.parametrize( + "positions, expected", + [ + ({}, TrajectoryExecutionStatus.START_STATE_UNAVAILABLE), + ({"joint": -1.0}, TrajectoryExecutionStatus.START_STATE_MISMATCH), + ], +) +def test_completed_trajectory_does_not_bypass_start_validation(positions, expected): + task = JointTrajectoryTask(JointTrajectoryTaskConfig(joint_names=["joint"])) + initial = JointTrajectory(joint_names=["joint"], points=[TrajectoryPoint(positions=[0.1])]) + assert task.execute(initial, {}).status is TrajectoryExecutionStatus.ACCEPTED + task.compute( + CoordinatorState( + joints=JointStateSnapshot(joint_positions={"joint": 0.0}), t_now=0.1, dt=0.1 + ) + ) + trajectory = JointTrajectory( + joint_names=["joint"], + points=[ + TrajectoryPoint(positions=[0.1]), + TrajectoryPoint(positions=[1.0], time_from_start=2.0), + ], + ) + assert task.execute(trajectory, positions).status is expected diff --git a/dimos/control/tasks/trajectory_task/trajectory_task.py b/dimos/control/tasks/trajectory_task/trajectory_task.py index e28cf428bf..7a20538788 100644 --- a/dimos/control/tasks/trajectory_task/trajectory_task.py +++ b/dimos/control/tasks/trajectory_task/trajectory_task.py @@ -58,7 +58,6 @@ def joint_trajectory_task( priority: int = 10, start_position_tolerance: float = 0.05, velocity_limits: Mapping[str, float] | None = None, - hold_position_when_idle: bool = False, ) -> TaskConfig: """Build the coordinator's single canonical joint-trajectory task.""" # The coordinator imports this module to recognize the canonical JTT. @@ -67,8 +66,6 @@ def joint_trajectory_task( params: dict[str, Any] = {"start_position_tolerance": start_position_tolerance} if velocity_limits is not None: params["velocity_limits"] = dict(velocity_limits) - if hold_position_when_idle: - params["hold_position_when_idle"] = True return TaskConfig( name=JOINT_TRAJECTORY_TASK_NAME, type="trajectory", @@ -142,8 +139,6 @@ class JointTrajectoryTaskConfig: position and the first trajectory point. velocity_limits: Optional positive velocity limit for every configured joint. Defaults to 1 rad/s per joint. - hold_position_when_idle: Keep emitting the last commanded position, - latching measured positions before the first trajectory. """ joint_names: Annotated[ @@ -157,7 +152,6 @@ class JointTrajectoryTaskConfig: allow_inf_nan=False, ) velocity_limits: dict[str, float] | None = None - hold_position_when_idle: bool = False @dataclass @@ -254,7 +248,7 @@ def on_joint_command(self, msg: JointState, t_now: float) -> bool: def is_active(self) -> bool: """Check if task should run this tick.""" - return self._config.hold_position_when_idle or self._state == TrajectoryState.EXECUTING + return self._state == TrajectoryState.EXECUTING def compute(self, state: CoordinatorState) -> JointCommandOutput | None: """Compute trajectory output for this tick. @@ -267,33 +261,10 @@ def compute(self, state: CoordinatorState) -> JointCommandOutput | None: Returns: JointCommandOutput with positions, or None if not executing """ - if self._config.hold_position_when_idle: - for joint_name in self._joint_names_list: - if joint_name in self._commanded_positions: - continue - measured = state.joints.get_position(joint_name) - if measured is not None and math.isfinite(measured): - self._commanded_positions[joint_name] = measured - if not self._motions: - if not self._config.hold_position_when_idle: - return None - held_names = [ - name for name in self._joint_names_list if name in self._commanded_positions - ] - if not held_names: - return None - return JointCommandOutput( - joint_names=held_names, - positions=[self._commanded_positions[name] for name in held_names], - mode=ControlMode.SERVO_POSITION, - ) + return None - output_names = ( - self._joint_names_list - if self._config.hold_position_when_idle - else [name for name in self._joint_names_list if name in self._motions] - ) + output_names = [name for name in self._joint_names_list if name in self._motions] all_complete = bool(self._motions) for joint_name, (run, index) in list(self._motions.items()): if run.start_time is None: @@ -332,11 +303,19 @@ def compute(self, state: CoordinatorState) -> JointCommandOutput | None: emitted_names = [name for name in output_names if name in self._commanded_positions] if not emitted_names: return None - return JointCommandOutput( + output = JointCommandOutput( joint_names=emitted_names, positions=[self._commanded_positions[name] for name in emitted_names], mode=ControlMode.SERVO_POSITION, ) + # Emit the final command, then forget joints we no longer control. + # Another task may move them before the next execution. + self._commanded_positions = { + name: position + for name, position in self._commanded_positions.items() + if name in self._motions + } + return output def on_preempted(self, by_task: str, joints: frozenset[str]) -> None: """Handle preemption by higher-priority task. @@ -587,7 +566,6 @@ class JointTrajectoryTaskParams(BaseConfig): allow_inf_nan=False, ) velocity_limits: dict[str, float] | None = None - hold_position_when_idle: bool = False def create_task(cfg: Any, hardware: Any) -> JointTrajectoryTask: @@ -602,6 +580,5 @@ def create_task(cfg: Any, hardware: Any) -> JointTrajectoryTask: priority=cfg.priority, start_position_tolerance=params.start_position_tolerance, velocity_limits=params.velocity_limits, - hold_position_when_idle=params.hold_position_when_idle, ), ) diff --git a/dimos/control/test_control.py b/dimos/control/test_control.py index 0048369935..bff9528cf8 100644 --- a/dimos/control/test_control.py +++ b/dimos/control/test_control.py @@ -578,55 +578,6 @@ def test_initial_state(self, trajectory_task): assert not trajectory_task.is_active() assert trajectory_task.get_state() == TrajectoryState.IDLE - def test_idle_hold_latches_measured_positions(self): - task = JointTrajectoryTask( - JointTrajectoryTaskConfig( - joint_names=["arm/joint1", "arm/joint2"], - hold_position_when_idle=True, - ) - ) - state = JointStateSnapshot(joint_positions={"arm/joint1": 0.25, "arm/joint2": -0.5}) - - output = task.compute(CoordinatorState(joints=state, t_now=1.0, dt=0.1)) - - assert task.is_active() - assert output is not None - assert output.joint_names == ["arm/joint1", "arm/joint2"] - assert output.positions == [0.25, -0.5] - - def test_idle_hold_retains_final_target_after_trajectory(self): - task = JointTrajectoryTask( - JointTrajectoryTaskConfig( - joint_names=["arm/joint1", "arm/joint2"], - start_position_tolerance=2.0, - velocity_limits={"arm/joint1": 10.0, "arm/joint2": 10.0}, - hold_position_when_idle=True, - ) - ) - trajectory = JointTrajectory( - joint_names=["arm/joint1"], - points=[ - TrajectoryPoint( - positions=[1.0], - velocities=[0.0], - time_from_start=0.0, - ) - ], - ) - state = JointStateSnapshot(joint_positions={"arm/joint1": 0.0, "arm/joint2": -0.5}) - assert ( - task.execute(trajectory, {"arm/joint1": 0.0}).status - is TrajectoryExecutionStatus.ACCEPTED - ) - - completed = task.compute(CoordinatorState(joints=state, t_now=1.0, dt=0.1)) - held = task.compute(CoordinatorState(joints=state, t_now=1.1, dt=0.1)) - - assert completed is not None - assert completed.positions == [1.0, -0.5] - assert held is not None - assert held.positions == [1.0, -0.5] - def test_claim(self, trajectory_task): claim = trajectory_task.claim() assert claim.priority == 10 diff --git a/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_groot_wbc.py b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_groot_wbc.py index 77218f927e..ce25f35b27 100644 --- a/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_groot_wbc.py +++ b/dimos/robot/unitree/g1/blueprints/basic/unitree_g1_groot_wbc.py @@ -283,12 +283,6 @@ def _precomposed_g1_scene(package: ScenePackage) -> Path | None: _default_ramp_seconds = 0.0 _decimation: int | None = 1 _n_workers = 2 # sim: keep the default worker count - _arm_holder = joint_trajectory_task( - g1_arms, - priority=10, - velocity_limits={name: 1.0 for name in g1_arms}, - hold_position_when_idle=True, - ) _mapper = VoxelGridMapper.blueprint(emit_every=1) _nav_stack = autoconnect( _mapper, @@ -327,12 +321,6 @@ def _precomposed_g1_scene(package: ScenePackage) -> Path | None: _decimation = 2 # 100 Hz tick / 2 = 50 Hz policy (training + sim rate). # One process per heavy module; fewer workers starve the Rerun bridge. _n_workers = 10 - _arm_holder = joint_trajectory_task( - g1_arms, - priority=10, - velocity_limits={name: 1.0 for name in g1_arms}, - hold_position_when_idle=True, - ) # Same nav middle as unitree-g1-nav-simple, fed by Point-LIO from the # MID-360, executed through the coordinator's twist_command. _nav_stack = autoconnect( @@ -362,6 +350,13 @@ def _precomposed_g1_scene(package: ScenePackage) -> Path | None: _nav_remappings = [] +_arm_trajectory_task = joint_trajectory_task( + g1_arms, + priority=10, + velocity_limits={name: 1.0 for name in g1_arms}, +) + + def _g1_groot_rerun_blueprint() -> Any: import rerun as rr import rerun.blueprint as rrb @@ -528,7 +523,7 @@ def _viewer() -> Any: "decimation": _decimation, }, ), - _arm_holder, + _arm_trajectory_task, # Shared bimanual Quest task with G1-only model and objective tuning. TaskConfig( name="teleop_g1",