diff --git a/docker-compose.yaml b/docker-compose.yaml index 190697fbb..169728fde 100644 --- a/docker-compose.yaml +++ b/docker-compose.yaml @@ -16,6 +16,9 @@ services: # MOVEIT_INFERENCE_PORT that publishes the server's host port, so # changing the port is a one-variable edit. No other config reads it. - INFER_URL=http://127.0.0.1:${MOVEIT_INFERENCE_PORT:-8973}/infer + # What the Trainer records as `observation.state`; see + # src/vla_sim/script/joint_command_bridge.py. + - MOVEIT_PRO_TRAIN_JOINT_STATES_TOPIC=${MOVEIT_PRO_TRAIN_JOINT_STATES_TOPIC:-/observed_joint_states} # Starts the robot drivers. drivers: diff --git a/src/vla_sim/CMakeLists.txt b/src/vla_sim/CMakeLists.txt index 343ca4b10..788330951 100644 --- a/src/vla_sim/CMakeLists.txt +++ b/src/vla_sim/CMakeLists.txt @@ -18,6 +18,7 @@ install( install(PROGRAMS script/get_action_chunk_adapter.py + script/joint_command_bridge.py DESTINATION lib/${PROJECT_NAME} ) @@ -33,6 +34,11 @@ if(BUILD_TESTING) TIMEOUT 60 ) + ament_add_pytest_test(test_joint_command_bridge + test/test_joint_command_bridge.py + TIMEOUT 60 + ) + ament_lint_auto_find_test_dependencies() endif() diff --git a/src/vla_sim/README.md b/src/vla_sim/README.md index 028607b56..fd5131f90 100644 --- a/src/vla_sim/README.md +++ b/src/vla_sim/README.md @@ -15,4 +15,19 @@ be too slow to run at all. A smaller model such as SmolVLA might be the better fit there. AMD GPUs are not passed through yet, so those machines run inference on the CPU as well. +## Collecting demonstrations + +The same scene also records the demonstrations a replacement policy is trained +on, one Pro dataset per prompt. + +`Run Cube-Stack Oracle` performs one stack from a scripted oracle, with no +recording, which is the quickest way to see whether a change to the scene or the +planner still produces a clean demonstration. `Collect Cube-Stack Demonstration` +records one. The six `Record Cube-Stack On ` objectives each +sweep the 60 training layouts drawn for their prompt, producing one dataset per +prompt. + +[`description/mujoco/keyframes.xml`](description/mujoco/keyframes.xml) holds the +layouts: 360 `train_*` and 150 `eval_*`. + For detailed documentation see: [MoveIt Pro Documentation](https://docs.picknik.ai/) diff --git a/src/vla_sim/config/config.yaml b/src/vla_sim/config/config.yaml index 4853da7a4..63d9fb065 100644 --- a/src/vla_sim/config/config.yaml +++ b/src/vla_sim/config/config.yaml @@ -16,6 +16,11 @@ hardware: # [Optional, default=True] launch_robot_state_publisher: True + # Hosts the joint command bridge, on both the dev and runtime paths. + additional_agent_launch_file: + package: "vla_sim" + path: "launch/simulated_extras.launch.py" + # Parameters used to configure the robot description through XACRO. # A URDF and SRDF are both required. # [Required] @@ -115,6 +120,8 @@ objectives: - "moveit_pro::behaviors::VisionBehaviorsLoader" - "moveit_pro::behaviors::ConverterBehaviorsLoader" - "moveit_pro::behaviors::MujocoBehaviorsLoader" + vla_sim: + - "vla_sim_behaviors::VlaSimBehaviorsLoader" # Specify source folder for objectives # [Required] objective_library_paths: diff --git a/src/vla_sim/description/mujoco/cube_stack_scene.xml b/src/vla_sim/description/mujoco/cube_stack_scene.xml index 8ceea1d5b..1b3a8ded2 100644 --- a/src/vla_sim/description/mujoco/cube_stack_scene.xml +++ b/src/vla_sim/description/mujoco/cube_stack_scene.xml @@ -1,4 +1,4 @@ - + - - - + + diff --git a/src/vla_sim/description/mujoco/keyframes.xml b/src/vla_sim/description/mujoco/keyframes.xml new file mode 100644 index 000000000..c0a2e3278 --- /dev/null +++ b/src/vla_sim/description/mujoco/keyframes.xml @@ -0,0 +1,3083 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/src/vla_sim/launch/simulated_extras.launch.py b/src/vla_sim/launch/simulated_extras.launch.py new file mode 100644 index 000000000..814693025 --- /dev/null +++ b/src/vla_sim/launch/simulated_extras.launch.py @@ -0,0 +1,45 @@ +#!/usr/bin/env python3 + +# Copyright 2026 PickNik Inc. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# * Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. +# +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# +# * Neither the name of the PickNik Inc. nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +# POSSIBILITY OF SUCH DAMAGE. + +from launch import LaunchDescription +from launch_ros.actions import Node + + +def generate_launch_description(): + return LaunchDescription( + [ + Node( + package="vla_sim", + executable="joint_command_bridge.py", + name="joint_command_bridge", + output="log", + ), + ] + ) diff --git a/src/vla_sim/objectives/collect_cube_stack_demo.xml b/src/vla_sim/objectives/collect_cube_stack_demo.xml new file mode 100644 index 000000000..f957b2d70 --- /dev/null +++ b/src/vla_sim/objectives/collect_cube_stack_demo.xml @@ -0,0 +1,32 @@ + + + + + + + + + + + + + + + + + + diff --git a/src/vla_sim/objectives/command_cube_stack_gripper.xml b/src/vla_sim/objectives/command_cube_stack_gripper.xml new file mode 100644 index 000000000..70d08d9dd --- /dev/null +++ b/src/vla_sim/objectives/command_cube_stack_gripper.xml @@ -0,0 +1,42 @@ + + + + + + + + + + + + + + + + + + + + diff --git a/src/vla_sim/objectives/execute_cube_stack_oracle.xml b/src/vla_sim/objectives/execute_cube_stack_oracle.xml new file mode 100644 index 000000000..61ec78221 --- /dev/null +++ b/src/vla_sim/objectives/execute_cube_stack_oracle.xml @@ -0,0 +1,69 @@ + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/src/vla_sim/objectives/move_along_cube_stack_keyposes.xml b/src/vla_sim/objectives/move_along_cube_stack_keyposes.xml new file mode 100644 index 000000000..8abc1f636 --- /dev/null +++ b/src/vla_sim/objectives/move_along_cube_stack_keyposes.xml @@ -0,0 +1,89 @@ + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/src/vla_sim/objectives/prepare_cube_stack_scene.xml b/src/vla_sim/objectives/prepare_cube_stack_scene.xml new file mode 100644 index 000000000..62124bf33 --- /dev/null +++ b/src/vla_sim/objectives/prepare_cube_stack_scene.xml @@ -0,0 +1,54 @@ + + + + + + + + + + + + + + + + + + + + + diff --git a/src/vla_sim/objectives/record_cube_stack_blue_on_green.xml b/src/vla_sim/objectives/record_cube_stack_blue_on_green.xml new file mode 100644 index 000000000..cad7aaed9 --- /dev/null +++ b/src/vla_sim/objectives/record_cube_stack_blue_on_green.xml @@ -0,0 +1,37 @@ + + + + + + + + + + + + + + + + + + diff --git a/src/vla_sim/objectives/record_cube_stack_blue_on_red.xml b/src/vla_sim/objectives/record_cube_stack_blue_on_red.xml new file mode 100644 index 000000000..ba680db32 --- /dev/null +++ b/src/vla_sim/objectives/record_cube_stack_blue_on_red.xml @@ -0,0 +1,37 @@ + + + + + + + + + + + + + + + + + + diff --git a/src/vla_sim/objectives/record_cube_stack_episode.xml b/src/vla_sim/objectives/record_cube_stack_episode.xml new file mode 100644 index 000000000..c5479fa10 --- /dev/null +++ b/src/vla_sim/objectives/record_cube_stack_episode.xml @@ -0,0 +1,56 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/src/vla_sim/objectives/record_cube_stack_green_on_blue.xml b/src/vla_sim/objectives/record_cube_stack_green_on_blue.xml new file mode 100644 index 000000000..d31df43b0 --- /dev/null +++ b/src/vla_sim/objectives/record_cube_stack_green_on_blue.xml @@ -0,0 +1,37 @@ + + + + + + + + + + + + + + + + + + diff --git a/src/vla_sim/objectives/record_cube_stack_green_on_red.xml b/src/vla_sim/objectives/record_cube_stack_green_on_red.xml new file mode 100644 index 000000000..b0460972f --- /dev/null +++ b/src/vla_sim/objectives/record_cube_stack_green_on_red.xml @@ -0,0 +1,37 @@ + + + + + + + + + + + + + + + + + + diff --git a/src/vla_sim/objectives/record_cube_stack_red_on_blue.xml b/src/vla_sim/objectives/record_cube_stack_red_on_blue.xml new file mode 100644 index 000000000..f7036e2f6 --- /dev/null +++ b/src/vla_sim/objectives/record_cube_stack_red_on_blue.xml @@ -0,0 +1,37 @@ + + + + + + + + + + + + + + + + + + diff --git a/src/vla_sim/objectives/record_cube_stack_red_on_green.xml b/src/vla_sim/objectives/record_cube_stack_red_on_green.xml new file mode 100644 index 000000000..41aabbde2 --- /dev/null +++ b/src/vla_sim/objectives/record_cube_stack_red_on_green.xml @@ -0,0 +1,37 @@ + + + + + + + + + + + + + + + + + + diff --git a/src/vla_sim/objectives/run_cube_stack_oracle.xml b/src/vla_sim/objectives/run_cube_stack_oracle.xml new file mode 100644 index 000000000..d1d55eb43 --- /dev/null +++ b/src/vla_sim/objectives/run_cube_stack_oracle.xml @@ -0,0 +1,30 @@ + + + + + + + + + + + + + + + + + diff --git a/src/vla_sim/package.xml b/src/vla_sim/package.xml index 36c2ebfc5..e1348e18f 100644 --- a/src/vla_sim/package.xml +++ b/src/vla_sim/package.xml @@ -22,6 +22,7 @@ picknik_mujoco_ros robotiq_description velocity_force_controller + vla_sim_behaviors moveit_pro_ml_msgs python3-numpy diff --git a/src/vla_sim/script/joint_command_bridge.py b/src/vla_sim/script/joint_command_bridge.py new file mode 100755 index 000000000..7f762c3f9 --- /dev/null +++ b/src/vla_sim/script/joint_command_bridge.py @@ -0,0 +1,286 @@ +#!/usr/bin/env python3 + +# Copyright 2026 PickNik Inc. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# * Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. +# +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# +# * Neither the name of the PickNik Inc. nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +# POSSIBILITY OF SUCH DAMAGE. + +"""Publish the training `action` and `observation.state` streams. + +Forge labels `action` from a joint command topic, which no MoveIt Pro Behavior +publishes. This node assembles the trajectory controller's reference setpoint +with the latched gripper command and republishes both as a +`sensor_msgs/JointState`. + +It publishes the same eight joints as `observation.state`, which is the width +the policies are trained on. Point MOVEIT_PRO_TRAIN_JOINT_STATES_TOPIC at that +topic; the default `/joint_states` carries all fifteen, eight of them passive +Robotiq linkage. + + ros2 run vla_sim joint_command_bridge.py \\ + --ros-args -p publish_rate:=10.0 +""" + +from __future__ import annotations + +import rclpy +from control_msgs.msg import JointTrajectoryControllerState +from rclpy.node import Node +from sensor_msgs.msg import JointState + +# joint_trajectory_admittance_controller publishes ~/controller_state only while +# a goal executes, so a reference older than this means no trajectory is running +# and the arm should be labelled with the pose it is holding. +DEFAULT_REFERENCE_TIMEOUT_S = 0.5 + +# The arm joints, in the order the trajectory controller reports them. The layout +# is fixed here so both streams publish from the first tick; the controller stays +# silent until a trajectory runs, several frames into an episode. +DEFAULT_ARM_JOINT_NAMES = [f"joint_{index}" for index in range(1, 8)] + + +def is_reference_fresh(now: float, stamp: float, timeout: float) -> bool: + return (now - stamp) < timeout + + +def assemble_joint_command( + joint_order: list[str], + gripper_joint: str | None, + reference_names: list[str] | None, + reference_positions: list[float] | None, + reference_fresh: bool, + measured: dict[str, float], + gripper_cmd: float | None, + gripper_fallback: float, +) -> list[float] | None: + """One position per name in ``joint_order``, or None if a joint has no source. + + Arm joints take the controller's setpoint while it is fresh, otherwise their + measured position. The gripper takes the latched command: its + ``GripperCommand`` goal is not observable on any topic, and it shares the + arm's radian scale through the ``split`` tendon in + ``description/mujoco/gen3_7dof.xml``. + """ + reference_map: dict[str, float] = {} + if ( + reference_fresh + and reference_names + and reference_positions + and len(reference_names) == len(reference_positions) + ): + reference_map = dict(zip(reference_names, reference_positions)) + + positions: list[float] = [] + for joint in joint_order: + if gripper_joint is not None and joint == gripper_joint: + positions.append( + float(gripper_cmd) + if gripper_cmd is not None + else float(gripper_fallback) + ) + elif joint in reference_map: + positions.append(float(reference_map[joint])) + elif joint in measured: + positions.append(float(measured[joint])) + else: + return None + return positions + + +class JointCommandBridge(Node): + def __init__(self) -> None: + super().__init__("joint_command_bridge") + + self._controller_state_topic = self._param( + "controller_state_topic", + "/joint_trajectory_admittance_controller/controller_state", + ) + self._joint_states_topic = self._param("joint_states_topic", "/joint_states") + self._joint_command_topic = self._param( + "joint_command_topic", "/joint_commands" + ) + self._observation_state_topic = self._param( + "observation_state_topic", "/observed_joint_states" + ) + self._arm_joint_names = list( + self._param("arm_joint_names", DEFAULT_ARM_JOINT_NAMES) + ) + self._gripper_joint = self._param( + "gripper_joint_name", "robotiq_85_left_knuckle_joint" + ) + self._gripper_fallback = float(self._param("gripper_rest_position", 0.0)) + # float(): keeps the declared type DOUBLE, which is what SetROS2Parameter + # sends at collection time. + self.declare_parameter("gripper_command_position", self._gripper_fallback) + # Forge resamples every stream onto the dataset's 1/fps grid. + self._publish_rate = float(self._param("publish_rate", 10.0)) + if self._publish_rate <= 0.0: + raise ValueError(f"publish_rate must be > 0, got {self._publish_rate}") + self._reference_timeout = float( + self._param("reference_timeout", DEFAULT_REFERENCE_TIMEOUT_S) + ) + + # One tuple, so names and positions are always read from the same message. + self._reference: tuple[list[str], list[float]] | None = None + self._reference_stamp = 0.0 + self._measured: dict[str, float] = {} + self._measured_stamp = 0.0 + # ExecutePolicy's layout: the planning group's order, gripper last. + self._joint_order = self._arm_joint_names + ( + [self._gripper_joint] if self._gripper_joint else [] + ) + + self._publisher = self.create_publisher( + JointState, self._joint_command_topic, 10 + ) + self._state_publisher = self.create_publisher( + JointState, self._observation_state_topic, 10 + ) + self.create_subscription( + JointTrajectoryControllerState, + self._controller_state_topic, + self._on_controller_state, + 10, + ) + self.create_subscription( + JointState, self._joint_states_topic, self._on_joint_states, 10 + ) + self.create_timer(1.0 / self._publish_rate, self._publish) + self.get_logger().info( + f"bridging {self._controller_state_topic} -> {self._joint_command_topic} " + f"and {self._joint_states_topic} -> {self._observation_state_topic} " + f"at {self._publish_rate:g} Hz for {self._joint_order}" + ) + + def _param(self, name: str, default): + return self.declare_parameter(name, default).value + + def _now(self) -> float: + return self.get_clock().now().nanoseconds * 1e-9 + + def _on_controller_state(self, msg: JointTrajectoryControllerState) -> None: + self._reference = (list(msg.joint_names), list(msg.reference.positions)) + self._reference_stamp = self._now() + if msg.joint_names and list(msg.joint_names) != self._arm_joint_names: + # The configured layout stands for the whole episode; report and keep + # the channel order stable. + self.get_logger().error( + f"{self._controller_state_topic} reports {list(msg.joint_names)}, not " + f"the configured arm_joint_names {self._arm_joint_names}; recorded " + "channels will not match the controller. Fix arm_joint_names.", + throttle_duration_sec=30.0, + ) + + def _on_joint_states(self, msg: JointState) -> None: + self._measured = dict(zip(msg.name, msg.position)) + self._measured_stamp = self._now() + + def _publish(self) -> None: + self._publish_action() + self._publish_observation() + + def _publish_action(self) -> None: + if ( + self._gripper_joint + and self._measured + and self._gripper_joint not in self._measured + ): + self.get_logger().warning( + f"gripper_joint_name '{self._gripper_joint}' is not in " + f"{self._joint_states_topic}; not publishing " + f"{self._joint_command_topic}. Check the configured joint name.", + throttle_duration_sec=30.0, + ) + return + + reference_names, reference_positions = self._reference or (None, None) + positions = assemble_joint_command( + joint_order=self._joint_order, + gripper_joint=self._gripper_joint or None, + reference_names=reference_names, + reference_positions=reference_positions, + reference_fresh=is_reference_fresh( + self._now(), self._reference_stamp, self._reference_timeout + ), + measured=self._measured, + gripper_cmd=self.get_parameter("gripper_command_position").value, + gripper_fallback=self._gripper_fallback, + ) + if positions is None: + self.get_logger().warning( + f"no setpoint or measured position for some joint in " + f"{self._joint_order}; not publishing {self._joint_command_topic}.", + throttle_duration_sec=10.0, + ) + return + self._publisher.publish(self._joint_state(positions)) + + def _publish_observation(self) -> None: + # Stale positions are dropped rather than restamped, so a frozen arm stays + # distinguishable from a still one. + if not is_reference_fresh( + self._now(), self._measured_stamp, self._reference_timeout + ): + self.get_logger().warning( + f"no {self._joint_states_topic} within {self._reference_timeout:g}s; " + f"not publishing {self._observation_state_topic}.", + throttle_duration_sec=10.0, + ) + return + missing = [j for j in self._joint_order if j not in self._measured] + if missing: + self.get_logger().warning( + f"{self._joint_states_topic} carries no {missing}; not publishing " + f"{self._observation_state_topic}.", + throttle_duration_sec=10.0, + ) + return + self._state_publisher.publish( + self._joint_state([self._measured[j] for j in self._joint_order]) + ) + + def _joint_state(self, positions: list[float]) -> JointState: + msg = JointState() + msg.header.stamp = self.get_clock().now().to_msg() + msg.name = list(self._joint_order) + msg.position = positions + return msg + + +def main() -> None: + rclpy.init() + node = JointCommandBridge() + try: + rclpy.spin(node) + except KeyboardInterrupt: + pass + finally: + node.destroy_node() + rclpy.try_shutdown() + + +if __name__ == "__main__": + main() diff --git a/src/vla_sim/test/test_joint_command_bridge.py b/src/vla_sim/test/test_joint_command_bridge.py new file mode 100644 index 000000000..b921c5de6 --- /dev/null +++ b/src/vla_sim/test/test_joint_command_bridge.py @@ -0,0 +1,190 @@ +#!/usr/bin/env python3 + +# Copyright 2026 PickNik Inc. +# All rights reserved. +# +# Unauthorized copying of this code base via any medium is strictly prohibited. +# Proprietary and confidential. + +"""Tests for the commanded-joint assembly that becomes the recorded `action`. + +Each case is a way the action could silently degrade into a copy of +`observation.state`, or into a constant, either of which trains a policy that +cannot move. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "script")) +from joint_command_bridge import ( # noqa: E402 + DEFAULT_REFERENCE_TIMEOUT_S, + assemble_joint_command, + is_reference_fresh, +) + +ARM = ["joint_a", "joint_b"] +GRIPPER = "gripper_joint" +ORDER = [*ARM, GRIPPER] + + +def test_returns_none_when_arm_joint_has_no_source() -> None: + """Publishing a short vector would misalign every downstream joint.""" + result = assemble_joint_command( + joint_order=ORDER, + gripper_joint=GRIPPER, + reference_names=None, + reference_positions=None, + reference_fresh=False, + measured={"joint_a": 0.1}, # joint_b absent + gripper_cmd=0.4, + gripper_fallback=0.0, + ) + assert result is None + + +def test_uses_controller_reference_for_arm_when_fresh() -> None: + """The setpoint leads the measured position, and that lead is the learning signal.""" + result = assemble_joint_command( + joint_order=ORDER, + gripper_joint=GRIPPER, + reference_names=["joint_a", "joint_b"], + reference_positions=[1.0, 2.0], + reference_fresh=True, + measured={"joint_a": 9.0, "joint_b": 9.0}, + gripper_cmd=0.6, + gripper_fallback=0.0, + ) + assert result == [1.0, 2.0, 0.6] + + +def test_falls_back_to_measured_when_reference_stale() -> None: + """A gripper dwell is labelled with the held pose, not the last goal's setpoint.""" + result = assemble_joint_command( + joint_order=ORDER, + gripper_joint=GRIPPER, + reference_names=["joint_a", "joint_b"], + reference_positions=[1.0, 2.0], + reference_fresh=False, + measured={"joint_a": 0.5, "joint_b": 0.6}, + gripper_cmd=0.6, + gripper_fallback=0.0, + ) + assert result == [0.5, 0.6, 0.6] + + +def test_gripper_uses_fallback_before_first_command() -> None: + result = assemble_joint_command( + joint_order=ORDER, + gripper_joint=GRIPPER, + reference_names=["joint_a", "joint_b"], + reference_positions=[1.0, 2.0], + reference_fresh=True, + measured={}, + gripper_cmd=None, + gripper_fallback=0.0, + ) + assert result == [1.0, 2.0, 0.0] + + +def test_gripper_zero_command_is_not_treated_as_unset() -> None: + """0.0 is the oracle's OPEN command and is falsy; the fallback differs to catch it.""" + result = assemble_joint_command( + joint_order=ORDER, + gripper_joint=GRIPPER, + reference_names=["joint_a", "joint_b"], + reference_positions=[1.0, 2.0], + reference_fresh=True, + measured={}, + gripper_cmd=0.0, + gripper_fallback=0.6, + ) + assert result == [1.0, 2.0, 0.0] + + +def test_mismatched_reference_lengths_are_ignored() -> None: + """Zipping them would bind joint_a to the wrong value and drop joint_b silently.""" + result = assemble_joint_command( + joint_order=ORDER, + gripper_joint=GRIPPER, + reference_names=["joint_a", "joint_b"], + reference_positions=[1.0], # length mismatch + reference_fresh=True, + measured={"joint_a": 0.5, "joint_b": 0.6}, + gripper_cmd=0.6, + gripper_fallback=0.0, + ) + assert result == [0.5, 0.6, 0.6] + + +def test_no_gripper_joint_builds_arm_only_vector() -> None: + result = assemble_joint_command( + joint_order=ARM, + gripper_joint=None, + reference_names=["joint_a", "joint_b"], + reference_positions=[1.0, 2.0], + reference_fresh=True, + measured={}, + gripper_cmd=None, + gripper_fallback=0.0, + ) + assert result == [1.0, 2.0] + + +def test_order_follows_joint_order_not_controller_order() -> None: + """A positional copy of the controller's vector would train on permuted labels.""" + result = assemble_joint_command( + joint_order=["joint_b", "joint_a", GRIPPER], + gripper_joint=GRIPPER, + reference_names=["joint_a", "joint_b"], + reference_positions=[1.0, 2.0], + reference_fresh=True, + measured={}, + gripper_cmd=0.6, + gripper_fallback=0.0, + ) + assert result == [2.0, 1.0, 0.6] + + +def test_gripper_never_takes_the_controller_reference() -> None: + """Checking the reference first would flatten the commanded open/close.""" + result = assemble_joint_command( + joint_order=ORDER, + gripper_joint=GRIPPER, + reference_names=["joint_a", "joint_b", GRIPPER], + reference_positions=[1.0, 2.0, 0.123], + reference_fresh=True, + measured={}, + gripper_cmd=0.6, + gripper_fallback=0.0, + ) + assert result == [1.0, 2.0, 0.6] + + +def test_reference_freshness_pins_the_timeout_on_both_sides() -> None: + """The timeout decides whether a dwell is labelled with the held pose.""" + timeout = DEFAULT_REFERENCE_TIMEOUT_S + assert is_reference_fresh(now=timeout - 1e-6, stamp=0.0, timeout=timeout) + assert not is_reference_fresh(now=timeout + 1e-6, stamp=0.0, timeout=timeout) + assert not is_reference_fresh( + now=timeout, stamp=0.0, timeout=timeout + ), "the boundary itself is stale" + + +def test_empty_joint_order_builds_an_empty_vector() -> None: + """The node never calls it this way, but the contract is one entry per name.""" + assert ( + assemble_joint_command( + joint_order=[], + gripper_joint=GRIPPER, + reference_names=["joint_a"], + reference_positions=[1.0], + reference_fresh=True, + measured={}, + gripper_cmd=0.6, + gripper_fallback=0.0, + ) + == [] + ) diff --git a/src/vla_sim_behaviors/CMakeLists.txt b/src/vla_sim_behaviors/CMakeLists.txt new file mode 100644 index 000000000..84834a0a3 --- /dev/null +++ b/src/vla_sim_behaviors/CMakeLists.txt @@ -0,0 +1,58 @@ +cmake_minimum_required(VERSION 3.22) +project(vla_sim_behaviors) + +find_package(moveit_pro_package REQUIRED) +moveit_pro_package() + +set(THIS_PACKAGE_INCLUDE_DEPENDS + control_msgs + geometry_msgs + moveit_pro_base + moveit_pro_behavior + moveit_pro_behavior_interface + moveit_studio_internal_msgs + nlohmann_json + pluginlib + rclcpp_action + sensor_msgs + tf2_eigen + trajectory_msgs) +foreach(package IN ITEMS ${THIS_PACKAGE_INCLUDE_DEPENDS}) + find_package(${package} REQUIRED) +endforeach() + +add_library( + vla_sim_behaviors + SHARED + src/compute_top_down_keyposes.cpp + src/plan_joint_spline_through_poses.cpp + src/send_gripper_command.cpp + src/wait_for_episode_start.cpp + src/register_behaviors.cpp) +target_include_directories( + vla_sim_behaviors + PUBLIC $ + $) +ament_target_dependencies(vla_sim_behaviors ${THIS_PACKAGE_INCLUDE_DEPENDS}) + +install( + TARGETS vla_sim_behaviors + EXPORT vla_sim_behaviorsTargets + ARCHIVE DESTINATION lib + LIBRARY DESTINATION lib + RUNTIME DESTINATION bin + INCLUDES + DESTINATION include) + +if(BUILD_TESTING) + find_package(ament_lint_auto REQUIRED) + moveit_pro_behavior_test(vla_sim_behaviors) + ament_lint_auto_find_test_dependencies() +endif() + +pluginlib_export_plugin_description_file( + moveit_pro_behavior_interface vla_sim_behaviors_plugin_description.xml) + +ament_export_targets(vla_sim_behaviorsTargets HAS_LIBRARY_TARGET) +ament_export_dependencies(${THIS_PACKAGE_INCLUDE_DEPENDS}) +ament_package() diff --git a/src/vla_sim_behaviors/include/vla_sim_behaviors/compute_top_down_keyposes.hpp b/src/vla_sim_behaviors/include/vla_sim_behaviors/compute_top_down_keyposes.hpp new file mode 100644 index 000000000..588c44c6f --- /dev/null +++ b/src/vla_sim_behaviors/include/vla_sim_behaviors/compute_top_down_keyposes.hpp @@ -0,0 +1,78 @@ +// Copyright 2026 PickNik Inc. +// All rights reserved. +// +// Unauthorized copying of this code base via any medium is strictly prohibited. +// Proprietary and confidential. + +#pragma once + +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace vla_sim_behaviors +{ +/** + * @brief Orientation of a top-down grasp whose jaw axis is spun to @p yaw about world +Z. + * + * The tip frame's +Z is the approach axis and +Y the jaw open/close axis, so pointing the + * approach axis at world -Z is Rx(pi). + */ +[[nodiscard]] Eigen::Quaterniond topDownGraspOrientation(double yaw); + +/** @brief Rotation about world +Z of @p orientation. */ +[[nodiscard]] double yawOf(const Eigen::Quaterniond& orientation); + +/** + * @brief The one of a cube's four symmetry-equivalent top-down yaws that @p cost scores lowest. + * + * A cube is 4-fold symmetric about its vertical axis, so all four grasps are visually identical + * and any of them stacks; which one the arm takes is purely a question of how far it has to + * move. @p cost is injected so the ranking can be scored by IK joint distance in the Behavior + * and by a plain metric in tests. + * + * @returns nullopt when @p cost scores no candidate, i.e. none is reachable. + */ +[[nodiscard]] std::optional chooseTopDownYawByCost(double cube_yaw, + const std::function(double)>& cost); + +/** @brief Sum of squared per-joint differences, the oracle's IK cost. */ +[[nodiscard]] double jointDistanceCost(const std::vector& from, const std::vector& to); + +/** + * @brief One top-down waypoint @p height above @p aim_pose, at @p orientation. + * + * @p held_object_offset is the held object's origin in the tip frame; it shifts the waypoint so + * the *object* lands on the aim point rather than the tip. Pass zero when the tip itself is the + * thing being positioned. + */ +[[nodiscard]] Eigen::Isometry3d topDownKeypose(const Eigen::Isometry3d& aim_pose, const Eigen::Quaterniond& orientation, + const Eigen::Vector3d& held_object_offset, double height); + +/** @brief topDownKeypose() at each of @p heights, all sharing @p orientation. */ +[[nodiscard]] std::vector computeTopDownKeyposes(const Eigen::Isometry3d& aim_pose, + const Eigen::Quaterniond& orientation, + const Eigen::Vector3d& held_object_offset, + const std::vector& heights); + +/** + * @brief Builds the Cartesian path for one segment of the cube stacking oracle. + */ +class ComputeTopDownKeyposes : public moveit_pro::behaviors::SharedResourcesNode +{ +public: + ComputeTopDownKeyposes(const std::string& name, const BT::NodeConfiguration& config, + const std::shared_ptr& shared_resources); + + [[nodiscard]] static BT::PortsList providedPorts(); + [[nodiscard]] static BT::KeyValueVector metadata(); + + BT::NodeStatus tick() override; +}; +} // namespace vla_sim_behaviors diff --git a/src/vla_sim_behaviors/include/vla_sim_behaviors/plan_joint_spline_through_poses.hpp b/src/vla_sim_behaviors/include/vla_sim_behaviors/plan_joint_spline_through_poses.hpp new file mode 100644 index 000000000..30132c091 --- /dev/null +++ b/src/vla_sim_behaviors/include/vla_sim_behaviors/plan_joint_spline_through_poses.hpp @@ -0,0 +1,101 @@ +// Copyright 2026 PickNik Inc. +// All rights reserved. +// +// Unauthorized copying of this code base via any medium is strictly prohibited. +// Proprietary and confidential. + +#pragma once + +#include + +#include +#include +#include +#include +#include + +namespace vla_sim_behaviors +{ +/** + * @brief Clamped cubic spline through joint-space knots. + * + * Passes through every knot exactly, is C1 continuous across them, and has zero first + * derivative at both ends. So the arm eases from rest, flows through the interior knots + * without stopping, and arrives at the last knot at rest. + * + * The parameter is dimensionless: divide the derivatives by the trajectory duration to + * get joint velocities and accelerations. + */ +class JointSpline +{ +public: + /** @throws std::invalid_argument unless there are >= 2 knots, all of one width, at + * strictly increasing @p parameters of the same count. */ + JointSpline(const std::vector& parameters, const std::vector& knots); + + [[nodiscard]] Eigen::VectorXd position(double s) const; + [[nodiscard]] Eigen::VectorXd velocity(double s) const; + [[nodiscard]] Eigen::VectorXd acceleration(double s) const; + + /** @brief Largest |d position / d parameter| over the spline, per joint. */ + [[nodiscard]] Eigen::VectorXd peakSpeed() const; + +private: + /** Where a parameter falls: the segment holding it, and the basis weights within it. */ + struct Segment + { + std::size_t index; + double width; + double from; + double to; + }; + + /** @p s saturated to the knot range, resolved against the segment containing it. */ + [[nodiscard]] Segment locate(double s) const; + + std::vector parameters_; + std::vector knots_; + std::vector moments_; +}; + +/** + * @brief How long one straight knot-to-knot segment should take. + * + * The slower of two budgets: covering @p cartesian_length at @p cartesian_speed, and + * holding every joint's peak speed under @p joint_velocity_cap. + */ +[[nodiscard]] double segmentDuration(double cartesian_length, const Eigen::VectorXd& joint_delta, + double cartesian_speed, const Eigen::VectorXd& joint_velocity_cap); + +/** + * @brief Knot parameters spanning [0, 1], each segment's share proportional to @p segment_durations. + * + * Giving a long or slow segment more of the parameter is what makes one spline over the + * whole chain respect a single speed budget instead of rushing its longest leg. + */ +[[nodiscard]] std::vector splineKnotParameters(const std::vector& segment_durations); + +/** + * @brief Duration to run @p spline over, as the slower of the same two budgets. + * + * Measured against the spline itself rather than the straight knot-to-knot deltas, since + * flowing through an interior knot overshoots what those deltas predict. + */ +[[nodiscard]] double splineDuration(const JointSpline& spline, double cartesian_length, double cartesian_speed, + const Eigen::VectorXd& joint_velocity_cap); + +/** + * @brief Plans one segment of the cube stacking oracle as a joint-space trajectory. + */ +class PlanJointSplineThroughPoses : public moveit_pro::behaviors::SharedResourcesNode +{ +public: + PlanJointSplineThroughPoses(const std::string& name, const BT::NodeConfiguration& config, + const std::shared_ptr& shared_resources); + + [[nodiscard]] static BT::PortsList providedPorts(); + [[nodiscard]] static BT::KeyValueVector metadata(); + + BT::NodeStatus tick() override; +}; +} // namespace vla_sim_behaviors diff --git a/src/vla_sim_behaviors/include/vla_sim_behaviors/send_gripper_command.hpp b/src/vla_sim_behaviors/include/vla_sim_behaviors/send_gripper_command.hpp new file mode 100644 index 000000000..926f16a04 --- /dev/null +++ b/src/vla_sim_behaviors/include/vla_sim_behaviors/send_gripper_command.hpp @@ -0,0 +1,37 @@ +// Copyright 2026 PickNik Inc. +// All rights reserved. +// +// Unauthorized copying of this code base via any medium is strictly prohibited. +// Proprietary and confidential. + +#pragma once + +#include +#include + +#include +#include +#include +#include + +namespace vla_sim_behaviors +{ +/** + * @brief Sends a gripper position goal and succeeds as soon as the server accepts it. + */ +class SendGripperCommand : public moveit_pro::behaviors::SharedResourcesNode +{ +public: + SendGripperCommand(const std::string& name, const BT::NodeConfiguration& config, + const std::shared_ptr& shared_resources); + + [[nodiscard]] static BT::PortsList providedPorts(); + [[nodiscard]] static BT::KeyValueVector metadata(); + + BT::NodeStatus tick() override; + +private: + rclcpp_action::Client::SharedPtr client_; + std::string action_name_; +}; +} // namespace vla_sim_behaviors diff --git a/src/vla_sim_behaviors/include/vla_sim_behaviors/wait_for_episode_start.hpp b/src/vla_sim_behaviors/include/vla_sim_behaviors/wait_for_episode_start.hpp new file mode 100644 index 000000000..f80780391 --- /dev/null +++ b/src/vla_sim_behaviors/include/vla_sim_behaviors/wait_for_episode_start.hpp @@ -0,0 +1,48 @@ +// Copyright 2026 PickNik Inc. +// All rights reserved. +// +// Unauthorized copying of this code base via any medium is strictly prohibited. +// Proprietary and confidential. + +#pragma once + +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace vla_sim_behaviors +{ +using GetActiveRecordingSrv = moveit_studio_internal_msgs::srv::GetActiveRecording; + +/** + * @brief Blocks until the active Trainer recording session has opened an episode. + */ +class WaitForEpisodeStart : public moveit_pro::behaviors::AsyncBehaviorBase +{ +public: + WaitForEpisodeStart(const std::string& name, const BT::NodeConfiguration& config, + const std::shared_ptr& shared_resources); + + [[nodiscard]] static BT::PortsList providedPorts(); + [[nodiscard]] static BT::KeyValueVector metadata(); + + tl::expected doWork() override; + tl::expected doHalt() override; + + std::shared_future>& getFuture() override + { + return future_; + } + +private: + std::shared_future> future_; + std::unique_ptr> client_; + std::atomic_bool halted_{ false }; +}; +} // namespace vla_sim_behaviors diff --git a/src/vla_sim_behaviors/package.xml b/src/vla_sim_behaviors/package.xml new file mode 100644 index 000000000..5c8d9c87a --- /dev/null +++ b/src/vla_sim_behaviors/package.xml @@ -0,0 +1,40 @@ + + + vla_sim_behaviors + 9.5.0 + + Behaviors for collecting cube stacking demonstrations with the vla_sim + configuration: the scripted stacking oracle's motion planning, a recordable + gripper command, and the recording start barrier. + + + MoveIt Pro Maintainer + + BSD-3-Clause + + ament_cmake + + moveit_pro_package + + control_msgs + geometry_msgs + moveit_pro_base + moveit_pro_behavior + moveit_pro_behavior_interface + moveit_studio_internal_msgs + nlohmann_json + pluginlib + rclcpp_action + sensor_msgs + trajectory_msgs + tf2_eigen + + ament_lint_auto + ament_cmake_gtest + ament_clang_format + ament_clang_tidy + + + ament_cmake + + diff --git a/src/vla_sim_behaviors/src/compute_top_down_keyposes.cpp b/src/vla_sim_behaviors/src/compute_top_down_keyposes.cpp new file mode 100644 index 000000000..515a498cc --- /dev/null +++ b/src/vla_sim_behaviors/src/compute_top_down_keyposes.cpp @@ -0,0 +1,278 @@ +// Copyright 2026 PickNik Inc. +// All rights reserved. +// +// Unauthorized copying of this code base via any medium is strictly prohibited. +// Proprietary and confidential. + +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace +{ +inline constexpr auto kDescriptionComputeTopDownKeyposes = R"( +

+ Builds the Cartesian path for one segment of the cube stacking oracle: + a top-down waypoint at each of heights above aim_pose, + all sharing one grasp orientation. +

+

+ The orientation points the tip's approach axis straight down, spun to whichever + of the cube's four symmetry-equivalent yaws costs the least arm motion from + seed_joint_state, scored by inverse kinematics. Set + held_object_offset to the grasped object's transform in the tip + frame to aim the object at the waypoints instead of the tip, which is what + places a carried cube on a target cube. +

+

+ Wire reuse_orientation from an earlier segment's + orientation once the gripper has closed, so the whole carry holds + the yaw the object was grasped at. +

+

+ Feed the scene from GetCurrentPlanningScene into + GetJointState, whose joint_state output is this + Behavior's seed_joint_state. Send path on to + PlanJointSplineThroughPoses. +

+ )"; + +constexpr auto kPortIDAimPose = "aim_pose"; +constexpr auto kPortIDSeedJointState = "seed_joint_state"; +constexpr auto kPortIDPlanningGroupName = "planning_group_name"; +constexpr auto kPortIDTipLink = "tip_link"; +constexpr auto kPortIDHeldObjectOffset = "held_object_offset"; +constexpr auto kPortIDHeights = "heights"; +constexpr auto kPortIDReuseOrientation = "reuse_orientation"; +constexpr auto kPortIDPath = "path"; +constexpr auto kPortIDOrientation = "orientation"; +} // namespace + +namespace vla_sim_behaviors +{ +Eigen::Quaterniond topDownGraspOrientation(double yaw) +{ + return Eigen::Quaterniond(Eigen::AngleAxisd(yaw, Eigen::Vector3d::UnitZ()) * + Eigen::AngleAxisd(M_PI, Eigen::Vector3d::UnitX())); +} + +double yawOf(const Eigen::Quaterniond& orientation) +{ + const Eigen::Matrix3d rotation = orientation.normalized().toRotationMatrix(); + return std::atan2(rotation(1, 0), rotation(0, 0)); +} + +std::optional chooseTopDownYawByCost(double cube_yaw, const std::function(double)>& cost) +{ + std::optional best_yaw; + std::optional best_cost; + for (int quarter_turns = 0; quarter_turns < 4; ++quarter_turns) + { + const double yaw = cube_yaw + quarter_turns * (M_PI / 2.0); + const std::optional candidate = cost(yaw); + if (candidate.has_value() && (!best_cost.has_value() || candidate.value() < best_cost.value())) + { + best_cost = candidate; + best_yaw = yaw; + } + } + return best_yaw; +} + +double jointDistanceCost(const std::vector& from, const std::vector& to) +{ + double cost = 0.0; + for (std::size_t i = 0; i < std::min(from.size(), to.size()); ++i) + { + const double difference = to[i] - from[i]; + cost += difference * difference; + } + return cost; +} + +Eigen::Isometry3d topDownKeypose(const Eigen::Isometry3d& aim_pose, const Eigen::Quaterniond& orientation, + const Eigen::Vector3d& held_object_offset, double height) +{ + // The offset is rigid in the tip frame, so rotating the wrist from the grasp orientation to + // this one carries the object with it. + const Eigen::Vector3d tip_to_object = orientation * held_object_offset; + Eigen::Isometry3d keypose(orientation); + keypose.translation() = aim_pose.translation() + Eigen::Vector3d(0.0, 0.0, height) - tip_to_object; + return keypose; +} + +std::vector computeTopDownKeyposes(const Eigen::Isometry3d& aim_pose, + const Eigen::Quaterniond& orientation, + const Eigen::Vector3d& held_object_offset, + const std::vector& heights) +{ + std::vector keyposes; + keyposes.reserve(heights.size()); + for (const double height : heights) + { + keyposes.push_back(topDownKeypose(aim_pose, orientation, held_object_offset, height)); + } + return keyposes; +} + +ComputeTopDownKeyposes::ComputeTopDownKeyposes( + const std::string& name, const BT::NodeConfiguration& config, + const std::shared_ptr& shared_resources) + : SharedResourcesNode(name, config, shared_resources) +{ +} + +BT::PortsList ComputeTopDownKeyposes::providedPorts() +{ + return { + BT::InputPort(kPortIDAimPose, "{aim_pose}", + "Pose the waypoints are stacked above, and whose yaw the grasp " + "aligns to. The cube to grasp or to stack onto; for a straight " + "retract, the current tip pose."), + BT::InputPort(kPortIDSeedJointState, "{seed_joint_state}", + "Arm's current joint positions, from GetJointState. Both the IK " + "seed and the pose the grasp candidates are scored against."), + BT::InputPort(kPortIDPlanningGroupName, "manipulator", + "SRDF joint group the grasp candidates are solved for."), + BT::InputPort(kPortIDTipLink, "grasp_link", "Link the waypoints position."), + BT::InputPort(kPortIDHeldObjectOffset, "", + "Grasped object's transform in the frame the waypoints " + "position: tip_link shifted by the planner's tip_offset. " + "Read it with GetLatestTransform using that frame as the " + "target. Leave unwired to position the frame itself."), + BT::InputPort>(kPortIDHeights, + "Heights above aim_pose, in meters, one waypoint each, semicolon separated."), + BT::InputPort(kPortIDReuseOrientation, "", + "Grasp orientation to reuse instead of choosing one. Wire this " + "from an earlier segment's orientation for any move that carries " + "a grasped object."), + BT::OutputPort>(kPortIDPath, "{path}", + "Cartesian path for PlanJointSplineThroughPoses, in " + "the frame of aim_pose."), + BT::OutputPort(kPortIDOrientation, "{orientation}", + "The grasp orientation these waypoints share."), + }; +} + +BT::KeyValueVector ComputeTopDownKeyposes::metadata() +{ + return { { moveit_pro::behaviors::kSubcategoryMetadataKey, "Cube Stacking" }, + { moveit_pro::behaviors::kDescriptionMetadataKey, kDescriptionComputeTopDownKeyposes } }; +} + +BT::NodeStatus ComputeTopDownKeyposes::tick() +{ + const auto ports = moveit_pro::behaviors::getRequiredInputs( + getInput(kPortIDAimPose), + getInput(kPortIDSeedJointState), getInput>(kPortIDHeights), + getInput(kPortIDPlanningGroupName), getInput(kPortIDTipLink)); + if (!ports.has_value()) + { + getBehaviorContext()->logger->publishFailureMessage( + name(), "Failed to get required values from input data ports: " + ports.error()); + return BT::NodeStatus::FAILURE; + } + const auto& [aim_pose_msg, seed_joint_state, heights, planning_group_name, tip_link] = ports.value(); + + if (heights.empty()) + { + getBehaviorContext()->logger->publishFailureMessage(name(), "heights is empty, so the path would have no " + "waypoints."); + return BT::NodeStatus::FAILURE; + } + + const auto& robot_model = getBehaviorContext()->robot_model; + const auto* joint_group = robot_model ? robot_model->getJointModelGroup(planning_group_name) : nullptr; + if (joint_group == nullptr) + { + getBehaviorContext()->logger->publishFailureMessage( + name(), fmt::format("No planning group '{}' in the robot model.", planning_group_name)); + return BT::NodeStatus::FAILURE; + } + + moveit_pro::base::RobotState seed_state(robot_model); + seed_state.setToDefaultValues(); + for (std::size_t i = 0; i < seed_joint_state.name.size() && i < seed_joint_state.position.size(); ++i) + { + if (robot_model->hasJointModel(seed_joint_state.name[i])) + { + seed_state.setJointPositions(seed_joint_state.name[i], { seed_joint_state.position[i] }); + } + } + seed_state.update(); + std::vector seed_positions; + seed_state.copyJointGroupPositions(joint_group, seed_positions); + + Eigen::Vector3d held_object_offset = Eigen::Vector3d::Zero(); + if (const auto offset = getInput(kPortIDHeldObjectOffset); offset.has_value()) + { + held_object_offset = tf2::transformToEigen(offset.value()).translation(); + } + + Eigen::Isometry3d aim_pose; + tf2::fromMsg(aim_pose_msg.pose, aim_pose); + + // Score each candidate where the arm arrives first, so the cost is the motion actually spent + // getting there rather than to the end of the segment. + const double approach_height = heights.front(); + const std::string ik_tip_link = tip_link; + const auto cost_of = [&](double yaw) -> std::optional { + const Eigen::Quaterniond orientation = topDownGraspOrientation(yaw); + const Eigen::Isometry3d keypose = topDownKeypose(aim_pose, orientation, held_object_offset, approach_height); + moveit_pro::base::RobotState candidate(seed_state); + if (!candidate.setFromIK(joint_group, keypose, ik_tip_link)) + { + return std::nullopt; + } + std::vector solution; + candidate.copyJointGroupPositions(joint_group, solution); + return jointDistanceCost(seed_positions, solution); + }; + + Eigen::Quaterniond orientation; + if (const auto reused = getInput(kPortIDReuseOrientation); reused.has_value()) + { + tf2::fromMsg(reused.value(), orientation); + } + else + { + const double cube_yaw = yawOf(Eigen::Quaterniond(aim_pose.rotation())); + const std::optional yaw = chooseTopDownYawByCost(cube_yaw, cost_of); + if (!yaw.has_value()) + { + getBehaviorContext()->logger->publishFailureMessage( + name(), "No top-down grasp of this pose is reachable: inverse kinematics failed for all four " + "symmetry-equivalent yaws."); + return BT::NodeStatus::FAILURE; + } + orientation = topDownGraspOrientation(yaw.value()); + } + + const std::vector keyposes = + computeTopDownKeyposes(aim_pose, orientation, held_object_offset, heights); + + std::vector path; + path.reserve(keyposes.size()); + for (const auto& keypose : keyposes) + { + geometry_msgs::msg::PoseStamped pose_msg; + pose_msg.header = aim_pose_msg.header; + pose_msg.pose = tf2::toMsg(keypose); + path.push_back(pose_msg); + } + setOutput(kPortIDPath, path); + setOutput(kPortIDOrientation, tf2::toMsg(orientation)); + + return BT::NodeStatus::SUCCESS; +} +} // namespace vla_sim_behaviors diff --git a/src/vla_sim_behaviors/src/plan_joint_spline_through_poses.cpp b/src/vla_sim_behaviors/src/plan_joint_spline_through_poses.cpp new file mode 100644 index 000000000..5c7bdbd0c --- /dev/null +++ b/src/vla_sim_behaviors/src/plan_joint_spline_through_poses.cpp @@ -0,0 +1,418 @@ +// Copyright 2026 PickNik Inc. +// All rights reserved. +// +// Unauthorized copying of this code base via any medium is strictly prohibited. +// Proprietary and confidential. + +#include + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace +{ +inline constexpr auto kDescriptionPlanJointSplineThroughPoses = R"( +

+ Plans a Cartesian path as a single joint-space trajectory: + each waypoint is solved for inverse kinematics once, warm-started from the + previous solution, and the resulting joint knots are joined by a clamped + cubic spline. The arm eases from rest, passes through every waypoint exactly + without stopping, and arrives at the last one at rest. +

+

+ Interpolating in joint space between same-branch solutions holds a redundant + arm's posture steady across the whole segment, which is what a demonstration + needs. +

+

+ The duration is the slower of covering the path at + cartesian_speed and holding every joint under its model velocity + limit scaled by joint_velocity_scale, so the tip holds a roughly + constant speed instead of a constant joint rate. +

+

+ Feed path from ComputeTopDownKeyposes, + seed_joint_state from GetJointState, and the + output to ExecuteTrajectory. +

+ )"; + +constexpr auto kPortIDPath = "path"; +constexpr auto kPortIDSeedJointState = "seed_joint_state"; +constexpr auto kPortIDPlanningGroupName = "planning_group_name"; +constexpr auto kPortIDTipLink = "tip_link"; +constexpr auto kPortIDTipOffset = "tip_offset"; +constexpr auto kPortIDCartesianSpeed = "cartesian_speed"; +constexpr auto kPortIDJointVelocityScale = "joint_velocity_scale"; +constexpr auto kPortIDSamplingRate = "sampling_rate"; +constexpr auto kPortIDJointTrajectory = "joint_trajectory_msg"; + +// Peak-to-mean speed ratio of a smootherstep profile. It sizes each segment's nominal duration, +// which in turn sets how much of the spline parameter that segment gets; the spline's true peak +// is measured afterwards. +constexpr double kNominalPeakSpeedRatio = 1.875; + +// A trajectory long enough to be worth executing, and a ceiling so an unreachable waypoint +// produces a failure rather than an arm that creeps for minutes. +constexpr double kMinimumDuration = 0.2; +constexpr double kMaximumDuration = 60.0; +} // namespace + +namespace vla_sim_behaviors +{ +JointSpline::JointSpline(const std::vector& parameters, const std::vector& knots) + : parameters_(parameters), knots_(knots) +{ + const std::size_t count = knots.size(); + if (count < 2) + { + throw std::invalid_argument(fmt::format("A spline needs at least 2 knots, got {}.", count)); + } + if (parameters.size() != count) + { + throw std::invalid_argument( + fmt::format("Got {} knots but {} parameters; they must correspond.", count, parameters.size())); + } + const Eigen::Index width = knots.front().size(); + for (std::size_t i = 0; i < count; ++i) + { + if (knots[i].size() != width) + { + throw std::invalid_argument( + fmt::format("Knot {} has width {}, but knot 0 has width {}.", i, knots[i].size(), width)); + } + if (i > 0 && !(parameters[i] > parameters[i - 1])) + { + throw std::invalid_argument(fmt::format("Parameters must strictly increase, but parameter {} is {} and {} is {}.", + i - 1, parameters[i - 1], i, parameters[i])); + } + } + + // Second derivatives at the knots, from the standard tridiagonal moment formulation with + // both ends clamped to zero first derivative. + Eigen::MatrixXd system = Eigen::MatrixXd::Zero(count, count); + Eigen::MatrixXd right_hand_side = Eigen::MatrixXd::Zero(count, width); + const auto span = [&](std::size_t i) { return parameters_[i + 1] - parameters_[i]; }; + const auto slope = [&](std::size_t i) -> Eigen::RowVectorXd { + return (knots_[i + 1] - knots_[i]).transpose() / span(i); + }; + + for (std::size_t i = 1; i + 1 < count; ++i) + { + system(i, i - 1) = span(i - 1); + system(i, i) = 2.0 * (span(i - 1) + span(i)); + system(i, i + 1) = span(i); + right_hand_side.row(i) = 6.0 * (slope(i) - slope(i - 1)); + } + system(0, 0) = 2.0 * span(0); + system(0, 1) = span(0); + right_hand_side.row(0) = 6.0 * slope(0); + system(count - 1, count - 1) = 2.0 * span(count - 2); + system(count - 1, count - 2) = span(count - 2); + right_hand_side.row(count - 1) = -6.0 * slope(count - 2); + + const Eigen::MatrixXd moments = system.colPivHouseholderQr().solve(right_hand_side); + moments_.reserve(count); + for (std::size_t i = 0; i < count; ++i) + { + moments_.push_back(moments.row(i).transpose()); + } +} + +JointSpline::Segment JointSpline::locate(double s) const +{ + // Saturating rather than extrapolating: a cubic continued past its last knot diverges fast. + s = std::clamp(s, parameters_.front(), parameters_.back()); + const auto upper = std::upper_bound(parameters_.begin(), parameters_.end(), s); + const auto found = static_cast(std::distance(parameters_.begin(), upper)); + const std::size_t index = std::clamp(found == 0 ? 0 : found - 1, 0, parameters_.size() - 2); + const double width = parameters_[index + 1] - parameters_[index]; + return { index, width, (parameters_[index + 1] - s) / width, (s - parameters_[index]) / width }; +} + +Eigen::VectorXd JointSpline::position(double s) const +{ + const auto [i, h, a, b] = locate(s); + return a * knots_[i] + b * knots_[i + 1] + + ((a * a * a - a) * moments_[i] + (b * b * b - b) * moments_[i + 1]) * (h * h / 6.0); +} + +Eigen::VectorXd JointSpline::velocity(double s) const +{ + const auto [i, h, a, b] = locate(s); + return (knots_[i + 1] - knots_[i]) / h + + ((1.0 - 3.0 * a * a) * moments_[i] + (3.0 * b * b - 1.0) * moments_[i + 1]) * (h / 6.0); +} + +Eigen::VectorXd JointSpline::acceleration(double s) const +{ + const auto [i, h, a, b] = locate(s); + return a * moments_[i] + b * moments_[i + 1]; +} + +Eigen::VectorXd JointSpline::peakSpeed() const +{ + Eigen::VectorXd peak = Eigen::VectorXd::Zero(knots_.front().size()); + for (std::size_t i = 0; i + 1 < parameters_.size(); ++i) + { + // Speed is quadratic within a segment, so its extremes are at the ends or where the + // acceleration crosses zero. + std::vector candidates = { parameters_[i], parameters_[i + 1] }; + for (Eigen::Index j = 0; j < peak.size(); ++j) + { + const double denominator = moments_[i + 1][j] - moments_[i][j]; + if (std::abs(denominator) > std::numeric_limits::epsilon()) + { + const double a = moments_[i + 1][j] / denominator; + if (a >= 0.0 && a <= 1.0) + { + candidates.push_back(parameters_[i + 1] - a * (parameters_[i + 1] - parameters_[i])); + } + } + } + for (const double s : candidates) + { + peak = peak.cwiseMax(velocity(s).cwiseAbs()); + } + } + return peak; +} + +double segmentDuration(double cartesian_length, const Eigen::VectorXd& joint_delta, double cartesian_speed, + const Eigen::VectorXd& joint_velocity_cap) +{ + const double cartesian = cartesian_speed > 0.0 ? cartesian_length / cartesian_speed : 0.0; + double joint = 0.0; + for (Eigen::Index j = 0; j < joint_delta.size() && j < joint_velocity_cap.size(); ++j) + { + if (joint_velocity_cap[j] > 0.0) + { + joint = std::max(joint, kNominalPeakSpeedRatio * std::abs(joint_delta[j]) / joint_velocity_cap[j]); + } + } + return std::max(cartesian, joint); +} + +std::vector splineKnotParameters(const std::vector& segment_durations) +{ + std::vector parameters{ 0.0 }; + parameters.reserve(segment_durations.size() + 1); + // Every segment gets a positive share, keeping the knot parameters strictly increasing. + const double floor = std::numeric_limits::epsilon(); + for (const double duration : segment_durations) + { + parameters.push_back(parameters.back() + std::max(duration, floor)); + } + const double total = parameters.back(); + for (double& parameter : parameters) + { + parameter /= total; + } + return parameters; +} + +double splineDuration(const JointSpline& spline, double cartesian_length, double cartesian_speed, + const Eigen::VectorXd& joint_velocity_cap) +{ + const double cartesian = cartesian_speed > 0.0 ? cartesian_length / cartesian_speed : 0.0; + const Eigen::VectorXd peak = spline.peakSpeed(); + double joint = 0.0; + for (Eigen::Index j = 0; j < peak.size() && j < joint_velocity_cap.size(); ++j) + { + if (joint_velocity_cap[j] > 0.0) + { + joint = std::max(joint, peak[j] / joint_velocity_cap[j]); + } + } + return std::clamp(std::max(cartesian, joint), kMinimumDuration, kMaximumDuration); +} + +PlanJointSplineThroughPoses::PlanJointSplineThroughPoses( + const std::string& name, const BT::NodeConfiguration& config, + const std::shared_ptr& shared_resources) + : SharedResourcesNode(name, config, shared_resources) +{ +} + +BT::PortsList PlanJointSplineThroughPoses::providedPorts() +{ + return { + BT::InputPort>(kPortIDPath, "{path}", + "Waypoints to pass through, in order."), + BT::InputPort(kPortIDSeedJointState, "{seed_joint_state}", + "Joint positions the trajectory starts from, from " + "GetJointState. Also seeds the first inverse kinematics " + "solve, which fixes the branch the whole path stays on."), + BT::InputPort(kPortIDPlanningGroupName, "manipulator", "SRDF joint group to plan for."), + BT::InputPort(kPortIDTipLink, "grasp_link", "Link the waypoints are solved for."), + BT::InputPort>(kPortIDTipOffset, "0;0;0", + "Translation from tip_link to the frame the waypoints actually position, " + "in the tip_link frame, semicolon separated."), + BT::InputPort(kPortIDCartesianSpeed, "0.065", "Tip speed budget along the path, in meters per second."), + BT::InputPort(kPortIDJointVelocityScale, "0.5", + "Fraction of each joint's model velocity limit the trajectory may reach. Binds only " + "where the path is short enough that cartesian_speed would exceed it."), + BT::InputPort(kPortIDSamplingRate, "100", "Output trajectory sampling rate, in Hz."), + BT::OutputPort(kPortIDJointTrajectory, "{joint_trajectory_msg}", + "Timed trajectory for ExecuteTrajectory."), + }; +} + +BT::KeyValueVector PlanJointSplineThroughPoses::metadata() +{ + return { { moveit_pro::behaviors::kSubcategoryMetadataKey, "Cube Stacking" }, + { moveit_pro::behaviors::kDescriptionMetadataKey, kDescriptionPlanJointSplineThroughPoses } }; +} + +BT::NodeStatus PlanJointSplineThroughPoses::tick() +{ + const auto ports = moveit_pro::behaviors::getRequiredInputs( + getInput>(kPortIDPath), + getInput(kPortIDSeedJointState), getInput(kPortIDPlanningGroupName), + getInput(kPortIDTipLink), getInput>(kPortIDTipOffset), + getInput(kPortIDCartesianSpeed), getInput(kPortIDJointVelocityScale), + getInput(kPortIDSamplingRate)); + if (!ports.has_value()) + { + getBehaviorContext()->logger->publishFailureMessage( + name(), "Failed to get required values from input data ports: " + ports.error()); + return BT::NodeStatus::FAILURE; + } + const auto& [path, seed_joint_state, planning_group_name, tip_link, tip_offset, cartesian_speed, joint_velocity_scale, + sampling_rate] = ports.value(); + + if (path.empty()) + { + getBehaviorContext()->logger->publishFailureMessage(name(), "path is empty, so there is nothing to plan."); + return BT::NodeStatus::FAILURE; + } + if (tip_offset.size() != 3) + { + getBehaviorContext()->logger->publishFailureMessage(name(), fmt::format("tip_offset needs 3 values, got {}.", + tip_offset.size())); + return BT::NodeStatus::FAILURE; + } + if (sampling_rate == 0) + { + getBehaviorContext()->logger->publishFailureMessage(name(), "sampling_rate must be positive."); + return BT::NodeStatus::FAILURE; + } + + const auto& robot_model = getBehaviorContext()->robot_model; + const auto* joint_group = robot_model ? robot_model->getJointModelGroup(planning_group_name) : nullptr; + if (joint_group == nullptr) + { + getBehaviorContext()->logger->publishFailureMessage( + name(), fmt::format("No planning group '{}' in the robot model.", planning_group_name)); + return BT::NodeStatus::FAILURE; + } + + moveit_pro::base::RobotState state(robot_model); + state.setToDefaultValues(); + for (std::size_t i = 0; i < seed_joint_state.name.size() && i < seed_joint_state.position.size(); ++i) + { + if (robot_model->hasJointModel(seed_joint_state.name[i])) + { + state.setJointPositions(seed_joint_state.name[i], { seed_joint_state.position[i] }); + } + } + state.update(); + + // The waypoints position a frame offset from tip_link, so the link itself has to arrive + // the same offset short of each one. + const Eigen::Translation3d offset(-tip_offset[0], -tip_offset[1], -tip_offset[2]); + + std::vector knots; + std::vector tip_positions; + { + std::vector positions; + state.copyJointGroupPositions(joint_group, positions); + knots.push_back(Eigen::Map(positions.data(), static_cast(positions.size()))); + tip_positions.push_back((state.getGlobalLinkTransform(tip_link) * offset.inverse()).translation()); + } + + for (std::size_t i = 0; i < path.size(); ++i) + { + Eigen::Isometry3d waypoint; + tf2::fromMsg(path[i].pose, waypoint); + if (!state.setFromIK(joint_group, waypoint * offset, tip_link)) + { + getBehaviorContext()->logger->publishFailureMessage( + name(), fmt::format("Inverse kinematics failed for waypoint {} of {}.", i + 1, path.size())); + return BT::NodeStatus::FAILURE; + } + std::vector positions; + state.copyJointGroupPositions(joint_group, positions); + knots.push_back(Eigen::Map(positions.data(), static_cast(positions.size()))); + tip_positions.push_back(waypoint.translation()); + } + + Eigen::VectorXd velocity_cap(knots.front().size()); + const auto& bounds = joint_group->getActiveJointModelsBounds(); + for (Eigen::Index j = 0; j < velocity_cap.size(); ++j) + { + const auto index = static_cast(j); + const bool bounded = index < bounds.size() && !bounds[index]->empty() && bounds[index]->front().velocity_bounded_; + velocity_cap[j] = + bounded ? bounds[index]->front().max_velocity_ * joint_velocity_scale : std::numeric_limits::infinity(); + } + + std::vector segment_durations; + double cartesian_length = 0.0; + segment_durations.reserve(knots.size() - 1); + for (std::size_t i = 0; i + 1 < knots.size(); ++i) + { + const double length = (tip_positions[i + 1] - tip_positions[i]).norm(); + cartesian_length += length; + segment_durations.push_back(segmentDuration(length, knots[i + 1] - knots[i], cartesian_speed, velocity_cap)); + } + + std::unique_ptr spline; + try + { + spline = std::make_unique(splineKnotParameters(segment_durations), knots); + } + catch (const std::invalid_argument& exception) + { + getBehaviorContext()->logger->publishFailureMessage(name(), fmt::format("Could not fit a spline through the " + "inverse kinematics solutions: {}", + exception.what())); + return BT::NodeStatus::FAILURE; + } + + const double duration = splineDuration(*spline, cartesian_length, cartesian_speed, velocity_cap); + const auto steps = static_cast(std::ceil(duration * sampling_rate)); + + trajectory_msgs::msg::JointTrajectory trajectory; + trajectory.header = path.front().header; + trajectory.joint_names = joint_group->getVariableNames(); + trajectory.points.reserve(steps + 1); + for (std::size_t step = 0; step <= steps; ++step) + { + const double s = static_cast(step) / static_cast(steps); + const Eigen::VectorXd position = spline->position(s); + const Eigen::VectorXd velocity = spline->velocity(s) / duration; + const Eigen::VectorXd acceleration = spline->acceleration(s) / (duration * duration); + + trajectory_msgs::msg::JointTrajectoryPoint point; + point.positions.assign(position.data(), position.data() + position.size()); + point.velocities.assign(velocity.data(), velocity.data() + velocity.size()); + point.accelerations.assign(acceleration.data(), acceleration.data() + acceleration.size()); + point.time_from_start = rclcpp::Duration::from_seconds(s * duration); + trajectory.points.push_back(std::move(point)); + } + setOutput(kPortIDJointTrajectory, trajectory); + + return BT::NodeStatus::SUCCESS; +} +} // namespace vla_sim_behaviors diff --git a/src/vla_sim_behaviors/src/register_behaviors.cpp b/src/vla_sim_behaviors/src/register_behaviors.cpp new file mode 100644 index 000000000..7d7718940 --- /dev/null +++ b/src/vla_sim_behaviors/src/register_behaviors.cpp @@ -0,0 +1,35 @@ +// Copyright 2026 PickNik Inc. +// All rights reserved. +// +// Unauthorized copying of this code base via any medium is strictly prohibited. +// Proprietary and confidential. + +#include +#include +#include + +#include +#include +#include +#include + +#include + +namespace vla_sim_behaviors +{ +class VlaSimBehaviorsLoader : public moveit_pro::behaviors::SharedResourcesNodeLoaderBase +{ +public: + void registerBehaviors(BT::BehaviorTreeFactory& factory, + const std::shared_ptr& shared_resources) override + { + moveit_pro::behaviors::registerBehavior(factory, "ComputeTopDownKeyposes", shared_resources); + moveit_pro::behaviors::registerBehavior(factory, "PlanJointSplineThroughPoses", + shared_resources); + moveit_pro::behaviors::registerBehavior(factory, "SendGripperCommand", shared_resources); + moveit_pro::behaviors::registerBehavior(factory, "WaitForEpisodeStart", shared_resources); + } +}; +} // namespace vla_sim_behaviors + +PLUGINLIB_EXPORT_CLASS(vla_sim_behaviors::VlaSimBehaviorsLoader, moveit_pro::behaviors::SharedResourcesNodeLoaderBase); diff --git a/src/vla_sim_behaviors/src/send_gripper_command.cpp b/src/vla_sim_behaviors/src/send_gripper_command.cpp new file mode 100644 index 000000000..b16dd2836 --- /dev/null +++ b/src/vla_sim_behaviors/src/send_gripper_command.cpp @@ -0,0 +1,106 @@ +// Copyright 2026 PickNik Inc. +// All rights reserved. +// +// Unauthorized copying of this code base via any medium is strictly prohibited. +// Proprietary and confidential. + +#include + +#include + +#include +#include +#include + +namespace +{ +inline constexpr auto kDescriptionSendGripperCommand = R"( +

+ Sends the gripper a position goal and succeeds as soon as the server accepts + it, without waiting for the jaws to arrive. +

+

+ This matches how ExecutePolicy drives the gripper at deploy time: + one goal per change of target, left to run while the arm keeps moving. Use it + to collect demonstrations of a grip the policy can reproduce. The goal stays + live afterwards, holding its commanded target against whatever the jaws have + closed onto. +

+

+ Pair with a WaitForDuration to hold the arm still while the jaws + travel. +

+ )"; + +constexpr auto kPortIDActionName = "gripper_command_action_name"; +constexpr auto kPortIDPosition = "position"; +constexpr auto kPortIDMaxEffort = "max_effort"; +constexpr auto kPortIDServerTimeout = "wait_for_server_timeout"; +} // namespace + +namespace vla_sim_behaviors +{ +SendGripperCommand::SendGripperCommand(const std::string& name, const BT::NodeConfiguration& config, + const std::shared_ptr& shared_resources) + : SharedResourcesNode(name, config, shared_resources) +{ +} + +BT::PortsList SendGripperCommand::providedPorts() +{ + return { + BT::InputPort(kPortIDActionName, "/robotiq_gripper_controller/gripper_cmd", + "GripperCommand action that actuates the gripper."), + BT::InputPort(kPortIDPosition, "Gripper joint target position."), + BT::InputPort(kPortIDMaxEffort, "0.0", "Effort ceiling; 0 leaves it to the controller."), + BT::InputPort(kPortIDServerTimeout, "3.0", "Seconds to wait for the action server to appear."), + }; +} + +BT::KeyValueVector SendGripperCommand::metadata() +{ + return { { moveit_pro::behaviors::kSubcategoryMetadataKey, "Cube Stacking" }, + { moveit_pro::behaviors::kDescriptionMetadataKey, kDescriptionSendGripperCommand } }; +} + +BT::NodeStatus SendGripperCommand::tick() +{ + const auto ports = + moveit_pro::behaviors::getRequiredInputs(getInput(kPortIDActionName), + getInput(kPortIDPosition), getInput(kPortIDMaxEffort), + getInput(kPortIDServerTimeout)); + if (!ports.has_value()) + { + getBehaviorContext()->logger->publishFailureMessage( + name(), "Failed to get required values from input data ports: " + ports.error()); + return BT::NodeStatus::FAILURE; + } + const auto& [action_name, position, max_effort, server_timeout] = ports.value(); + + // The client is kept between ticks: it outlives the goal request either way, and the + // shared node's executor is what delivers it. + if (client_ == nullptr || action_name_ != action_name) + { + client_ = + rclcpp_action::create_client(getBehaviorContext()->node, action_name); + action_name_ = action_name; + } + + const auto timeout = std::chrono::duration(server_timeout); + if (!client_->wait_for_action_server(std::chrono::duration_cast(timeout))) + { + getBehaviorContext()->logger->publishFailureMessage(name(), fmt::format("No GripperCommand action server on '{}'.", + action_name)); + return BT::NodeStatus::FAILURE; + } + + control_msgs::action::GripperCommand::Goal goal; + goal.command.position = position; + goal.command.max_effort = max_effort; + // The goal handle is dropped, since waiting on it is the blocking this Behavior exists to + // avoid. A rejected goal surfaces as a gripper that did not move. + client_->async_send_goal(goal); + + return BT::NodeStatus::SUCCESS; +} +} // namespace vla_sim_behaviors diff --git a/src/vla_sim_behaviors/src/wait_for_episode_start.cpp b/src/vla_sim_behaviors/src/wait_for_episode_start.cpp new file mode 100644 index 000000000..c94d5638f --- /dev/null +++ b/src/vla_sim_behaviors/src/wait_for_episode_start.cpp @@ -0,0 +1,147 @@ +// Copyright 2026 PickNik Inc. +// All rights reserved. +// +// Unauthorized copying of this code base via any medium is strictly prohibited. +// Proprietary and confidential. + +#include + +#include +#include + +#include +#include +#include +#include + +namespace +{ +inline constexpr auto kDescriptionWaitForEpisodeStart = R"( +

+ Blocks until the active Trainer recording session reports that it has opened + an episode, then succeeds. +

+

+ RecordEpisode returns once the recorder process is spawned, and + the episode opens seconds later when the Trainer publishes its start marker. + Conversion keeps only motion after that marker, so tick this between + RecordEpisode and the demonstrated motion. +

+

+ Polling the session state catches the transition whenever it lands. Pair it + with a short WaitForDuration to keep the arm still across the few + frames straddling the marker; the converter's idle trimming removes them. +

+

+ Fails when no episode opens within timeout, naming the last state + seen. +

+ )"; + +constexpr auto kPortIDServiceName = "service_name"; +constexpr auto kPortIDTimeout = "timeout"; + +constexpr auto kDefaultServiceName = "/trainer/active_recording"; + +// The RecordingState value that means the episode's start marker has been published. +constexpr auto kRecordingState = "recording"; + +constexpr std::chrono::milliseconds kPollPeriod{ 50 }; +constexpr std::chrono::duration kServerTimeout{ 5.0 }; +constexpr std::chrono::duration kResponseTimeout{ 5.0 }; + +/** @brief The session's state, or "" when no session is active or the payload has no state. */ +std::string stateOf(const std::string& session_json) +{ + if (session_json.empty()) + { + return ""; + } + const auto session = nlohmann::json::parse(session_json, nullptr, false); + if (session.is_discarded() || !session.contains("state") || !session["state"].is_string()) + { + return ""; + } + return session["state"].get(); +} +} // namespace + +namespace vla_sim_behaviors +{ +WaitForEpisodeStart::WaitForEpisodeStart(const std::string& name, const BT::NodeConfiguration& config, + const std::shared_ptr& shared_resources) + : AsyncBehaviorBase(name, config, shared_resources) + , client_(std::make_unique>(shared_resources)) +{ +} + +BT::PortsList WaitForEpisodeStart::providedPorts() +{ + return { + BT::InputPort(kPortIDServiceName, kDefaultServiceName, "Name of the Trainer active_recording service."), + BT::InputPort(kPortIDTimeout, "30.0", "Seconds to wait for an episode to open."), + }; +} + +BT::KeyValueVector WaitForEpisodeStart::metadata() +{ + return { { moveit_pro::behaviors::kSubcategoryMetadataKey, "Cube Stacking" }, + { moveit_pro::behaviors::kDescriptionMetadataKey, kDescriptionWaitForEpisodeStart } }; +} + +tl::expected WaitForEpisodeStart::doWork() +{ + const auto ports = moveit_pro::behaviors::getRequiredInputs(getInput(kPortIDServiceName), + getInput(kPortIDTimeout)); + if (!ports.has_value()) + { + return tl::make_unexpected("Failed to get required values from input data ports: " + ports.error()); + } + const auto& [service_name, timeout] = ports.value(); + + halted_ = false; + client_->initialize(service_name, kServerTimeout, kResponseTimeout); + if (!client_->waitForServiceServer()) + { + return tl::make_unexpected(fmt::format("No Trainer active_recording service on '{}'.", service_name)); + } + // Nothing here holds a goal, so a halt may interrupt the poll as soon as it arrives. + notifyCanHalt(); + + const auto deadline = std::chrono::steady_clock::now() + std::chrono::duration(timeout); + std::string last_state; + while (!halted_) + { + const auto response = client_->syncSendRequest(GetActiveRecordingSrv::Request{}); + if (!response.has_value()) + { + return tl::make_unexpected("Failed to read the active recording: " + response.error()); + } + if (!response.value().status.success) + { + return tl::make_unexpected("The Trainer refused to report the active recording: " + + response.value().status.error_message); + } + last_state = stateOf(response.value().session_json); + if (last_state == kRecordingState) + { + return true; + } + if (std::chrono::steady_clock::now() >= deadline) + { + return tl::make_unexpected( + fmt::format("No recording episode opened within {:g}s; the session's last state was '{}'.", timeout, + last_state.empty() ? "no active session" : last_state)); + } + std::this_thread::sleep_for(kPollPeriod); + } + return false; +} + +tl::expected WaitForEpisodeStart::doHalt() +{ + halted_ = true; + client_->cancelRequest(); + return {}; +} +} // namespace vla_sim_behaviors diff --git a/src/vla_sim_behaviors/test/CMakeLists.txt b/src/vla_sim_behaviors/test/CMakeLists.txt new file mode 100644 index 000000000..f0a9cea2e --- /dev/null +++ b/src/vla_sim_behaviors/test/CMakeLists.txt @@ -0,0 +1,24 @@ +find_package(ament_cmake_gtest REQUIRED) + +# rclcpp::init wedges on discovery when it shares a ROS domain with a live backend, +# so every test that brings up a node gets a domain of its own -- one each, since +# colcon runs them in parallel and they would otherwise discover each other. The +# domain is read while the rmw library loads, too early for the test's own main(). +ament_add_gtest(test_behavior_plugins test_behavior_plugins.cpp ENV "ROS_DOMAIN_ID=89") +ament_target_dependencies(test_behavior_plugins ${THIS_PACKAGE_INCLUDE_DEPENDS}) + +ament_add_gtest(test_compute_top_down_keyposes test_compute_top_down_keyposes.cpp) +target_link_libraries(test_compute_top_down_keyposes vla_sim_behaviors) +ament_target_dependencies(test_compute_top_down_keyposes ${THIS_PACKAGE_INCLUDE_DEPENDS}) + +ament_add_gtest(test_plan_joint_spline_through_poses test_plan_joint_spline_through_poses.cpp) +target_link_libraries(test_plan_joint_spline_through_poses vla_sim_behaviors) +ament_target_dependencies(test_plan_joint_spline_through_poses ${THIS_PACKAGE_INCLUDE_DEPENDS}) + +ament_add_gtest(test_send_gripper_command test_send_gripper_command.cpp ENV "ROS_DOMAIN_ID=90") +target_link_libraries(test_send_gripper_command vla_sim_behaviors) +ament_target_dependencies(test_send_gripper_command ${THIS_PACKAGE_INCLUDE_DEPENDS}) + +ament_add_gtest(test_wait_for_episode_start test_wait_for_episode_start.cpp ENV "ROS_DOMAIN_ID=91") +target_link_libraries(test_wait_for_episode_start vla_sim_behaviors) +ament_target_dependencies(test_wait_for_episode_start ${THIS_PACKAGE_INCLUDE_DEPENDS}) diff --git a/src/vla_sim_behaviors/test/test_behavior_plugins.cpp b/src/vla_sim_behaviors/test/test_behavior_plugins.cpp new file mode 100644 index 000000000..12a30c06a --- /dev/null +++ b/src/vla_sim_behaviors/test/test_behavior_plugins.cpp @@ -0,0 +1,46 @@ +// Copyright 2026 PickNik Inc. +// All rights reserved. +// +// Unauthorized copying of this code base via any medium is strictly prohibited. +// Proprietary and confidential. + +#include + +#include +#include +#include +#include + +/** + * @brief This test makes sure that the Behaviors provided in this package can be successfully registered and + * instantiated by the behavior tree factory. + */ +TEST(BehaviorTests, test_load_behavior_plugins) +{ + pluginlib::ClassLoader class_loader( + "moveit_pro_behavior_interface", "moveit_pro::behaviors::SharedResourcesNodeLoaderBase"); + + auto node = std::make_shared("test_node"); + auto shared_resources = std::make_shared(node); + + BT::BehaviorTreeFactory factory; + { + auto plugin_instance = class_loader.createUniqueInstance("vla_sim_behaviors::VlaSimBehaviorsLoader"); + ASSERT_NO_THROW(plugin_instance->registerBehaviors(factory, shared_resources)); + } + // Test that ClassLoader is able to find and instantiate each Behavior using the package's plugin description info. + for (const auto& behavior_name : + { "ComputeTopDownKeyposes", "PlanJointSplineThroughPoses", "SendGripperCommand", "WaitForEpisodeStart" }) + { + EXPECT_NO_THROW((void)factory.instantiateTreeNode("test_behavior_name", behavior_name, BT::NodeConfiguration())) + << "Behavior '" << behavior_name << "' is registered but could not be instantiated."; + } +} + +int main(int argc, char** argv) +{ + rclcpp::init(argc, argv); + + testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/src/vla_sim_behaviors/test/test_compute_top_down_keyposes.cpp b/src/vla_sim_behaviors/test/test_compute_top_down_keyposes.cpp new file mode 100644 index 000000000..da01de96a --- /dev/null +++ b/src/vla_sim_behaviors/test/test_compute_top_down_keyposes.cpp @@ -0,0 +1,229 @@ +// Copyright 2026 PickNik Inc. +// All rights reserved. +// +// Unauthorized copying of this code base via any medium is strictly prohibited. +// Proprietary and confidential. + +#include + +#include +#include +#include + +#include + +namespace +{ +using vla_sim_behaviors::chooseTopDownYawByCost; +using vla_sim_behaviors::computeTopDownKeyposes; +using vla_sim_behaviors::jointDistanceCost; +using vla_sim_behaviors::topDownGraspOrientation; +using vla_sim_behaviors::yawOf; + +constexpr double kEpsilon = 1e-9; +constexpr double kQuarterTurn = M_PI / 2.0; + +Eigen::Isometry3d makePose(const Eigen::Vector3d& position, double yaw) +{ + Eigen::Isometry3d pose(Eigen::AngleAxisd(yaw, Eigen::Vector3d::UnitZ())); + pose.translation() = position; + return pose; +} + +// The oracle's Q_TOPDOWN, in the (x, y, z, w) order Eigen's constructor takes last. +Eigen::Quaterniond oracleTopDown() +{ + return Eigen::Quaterniond(0.0, 1.0, 0.0, 0.0); +} +} // namespace + +TEST(TopDownGraspOrientation, ZeroYawMatchesTheOracleBaseOrientation) +{ + // mujoco_ik.py pins Q_TOPDOWN = (w, x, y, z) = (0, 1, 0, 0). A different sign convention or + // axis here would still point the jaws down but spin them 90 degrees off the cube. + EXPECT_NEAR(std::abs(topDownGraspOrientation(0.0).dot(oracleTopDown())), 1.0, kEpsilon); +} + +TEST(TopDownGraspOrientation, ApproachAxisPointsDown) +{ + for (const double yaw : { -2.0, 0.0, 0.7, 3.0 }) + { + const Eigen::Vector3d approach = topDownGraspOrientation(yaw) * Eigen::Vector3d::UnitZ(); + EXPECT_NEAR(approach.z(), -1.0, kEpsilon) << "yaw " << yaw; + } +} + +TEST(TopDownGraspOrientation, SpinsTheJawAxisByYaw) +{ + // The jaw axis is the tip's +Y. Spinning by yaw must rotate it in the horizontal plane, or the + // grasp does not line up with the cube's faces. + // Rx(pi) flips +Y to -Y, then Rz(pi/2) carries that onto +X. + const Eigen::Vector3d jaw = topDownGraspOrientation(kQuarterTurn) * Eigen::Vector3d::UnitY(); + EXPECT_NEAR(jaw.x(), 1.0, kEpsilon); + EXPECT_NEAR(jaw.z(), 0.0, kEpsilon); +} + +TEST(YawOf, RoundTripsTopDownGraspOrientation) +{ + for (const double yaw : { -1.5, -0.2, 0.0, 0.9 }) + { + EXPECT_NEAR(yawOf(topDownGraspOrientation(yaw)), yaw, kEpsilon) << "yaw " << yaw; + } +} + +TEST(JointDistanceCost, IsZeroForTheSamePose) +{ + const std::vector pose{ 0.1, -0.2, 0.3 }; + EXPECT_NEAR(jointDistanceCost(pose, pose), 0.0, kEpsilon); +} + +TEST(JointDistanceCost, SumsTheSquaredPerJointDifferences) +{ + // The oracle's score, so a joint that moves twice as far counts four times as much and one + // big wrist swing outweighs several small arm adjustments. + EXPECT_NEAR(jointDistanceCost({ 0.0, 0.0 }, { 3.0, 4.0 }), 25.0, kEpsilon); +} + +TEST(ChooseTopDownYawByCost, ReturnsACubeYawPlusAWholeNumberOfQuarterTurns) +{ + // Every candidate must be a symmetry of the cube. A yaw that is not one grasps a corner. + const double chosen = + chooseTopDownYawByCost(0.3, [](double yaw) { return std::optional(std::abs(yaw)); }).value(); + const double turns = (chosen - 0.3) / kQuarterTurn; + EXPECT_NEAR(turns, std::round(turns), kEpsilon); +} + +TEST(ChooseTopDownYawByCost, PicksTheCheapestCandidateRatherThanTheNearest) +{ + // The whole point of scoring by IK: on eval_0 the oracle takes a candidate 180 degrees from + // the wrist's current yaw because it holds joint_5 still. A nearest-yaw rule cannot do that. + const double cube_yaw = 0.0; + const auto cost = [](double yaw) -> std::optional { + // cheapest at two quarter turns, i.e. the candidate a half turn away + return std::abs(std::remainder(yaw - M_PI, 2.0 * M_PI)); + }; + EXPECT_NEAR(chooseTopDownYawByCost(cube_yaw, cost).value(), M_PI, kEpsilon); +} + +TEST(ChooseTopDownYawByCost, SkipsUnreachableCandidates) +{ + // IK fails on candidates that would put the wrist past a limit; those must not be chosen + // even when a reachable one scores worse. + const auto cost = [](double yaw) -> std::optional { + if (std::abs(std::remainder(yaw, 2.0 * M_PI)) < kEpsilon) + { + return 0.0; // cheapest, but pretend it is the only reachable one below + } + return std::nullopt; + }; + EXPECT_NEAR(chooseTopDownYawByCost(0.0, cost).value(), 0.0, kEpsilon); +} + +TEST(ChooseTopDownYawByCost, ReturnsNulloptWhenNothingIsReachable) +{ + // The segment must fail loudly rather than plan a path the arm cannot follow. + EXPECT_FALSE(chooseTopDownYawByCost(0.4, [](double) { return std::nullopt; }).has_value()); +} + +TEST(ChooseTopDownYawByCost, IsUnchangedByCubeYawsAQuarterTurnApart) +{ + // The cube's own yaw is only known modulo a quarter turn, so equivalent readings of the same + // physical cube must produce the same grasp. + const auto cost = [](double yaw) -> std::optional { return std::abs(std::remainder(yaw - 0.9, 2.0 * M_PI)); }; + const double base = chooseTopDownYawByCost(0.2, cost).value(); + EXPECT_NEAR(std::remainder(chooseTopDownYawByCost(0.2 + kQuarterTurn, cost).value() - base, 2.0 * M_PI), 0.0, + kEpsilon); + EXPECT_NEAR(std::remainder(chooseTopDownYawByCost(0.2 - kQuarterTurn, cost).value() - base, 2.0 * M_PI), 0.0, + kEpsilon); +} + +TEST(ComputeTopDownKeyposes, StacksOneWaypointPerHeightAboveTheAimPose) +{ + const auto keyposes = computeTopDownKeyposes(makePose({ 0.5, -0.1, 0.115 }, 0.0), topDownGraspOrientation(0.0), + Eigen::Vector3d::Zero(), { 0.12, 0.0 }); + + ASSERT_EQ(keyposes.size(), 2u); + EXPECT_NEAR(keyposes[0].translation().z(), 0.235, kEpsilon); + EXPECT_NEAR(keyposes[1].translation().z(), 0.115, kEpsilon); + for (const auto& keypose : keyposes) + { + EXPECT_NEAR(keypose.translation().x(), 0.5, kEpsilon); + EXPECT_NEAR(keypose.translation().y(), -0.1, kEpsilon); + } +} + +TEST(ComputeTopDownKeyposes, GivesEveryWaypointTheChosenOrientation) +{ + // The oracle descends straight down onto the cube. Re-deriving the orientation per waypoint + // would let the wrist rotate mid-descent and shear the grasp. + const Eigen::Quaterniond orientation = topDownGraspOrientation(0.4); + const auto keyposes = computeTopDownKeyposes(makePose({ 0.5, 0.0, 0.115 }, 0.4), orientation, Eigen::Vector3d::Zero(), + { 0.12, 0.06, 0.0 }); + + ASSERT_EQ(keyposes.size(), 3u); + for (const auto& keypose : keyposes) + { + EXPECT_NEAR(std::abs(Eigen::Quaterniond(keypose.rotation()).dot(orientation)), 1.0, kEpsilon); + } +} + +TEST(ComputeTopDownKeyposes, PlacesTheHeldObjectRatherThanTheTipWhenOffsetIsSet) +{ + // The place segment aims the carried cube at the target, so the tip must land offset by + // exactly the grip, rotated into the world. + const Eigen::Vector3d grip_offset(0.0, 0.0, 0.02); + const auto keyposes = + computeTopDownKeyposes(makePose({ 0.4, 0.2, 0.115 }, 0.0), topDownGraspOrientation(0.0), grip_offset, { 0.031 }); + + ASSERT_EQ(keyposes.size(), 1u); + // The tip's +Z points down, so an object 20 mm along +Z sits 20 mm below the tip: the tip goes + // that much higher for the object to land on the aim point. + EXPECT_NEAR(keyposes[0].translation().z(), 0.115 + 0.031 + 0.02, kEpsilon); + EXPECT_NEAR(keyposes[0].translation().x(), 0.4, kEpsilon); + EXPECT_NEAR(keyposes[0].translation().y(), 0.2, kEpsilon); +} + +TEST(ComputeTopDownKeyposes, RotatesALateralHeldObjectOffsetIntoTheWorld) +{ + // A grip that is off-center laterally must be corrected in the direction the wrist is actually + // pointing. Ignoring the rotation would put the cube down on the wrong side of the target. + const Eigen::Vector3d grip_offset(0.01, 0.0, 0.0); + const auto keyposes = computeTopDownKeyposes(makePose({ 0.4, 0.2, 0.115 }, kQuarterTurn), + topDownGraspOrientation(kQuarterTurn), grip_offset, { 0.0 }); + + ASSERT_EQ(keyposes.size(), 1u); + // Rz(pi/2) * Rx(pi) maps the tip's +X onto world +Y, so the tip shifts back along -Y. + EXPECT_NEAR(keyposes[0].translation().x(), 0.4, kEpsilon); + EXPECT_NEAR(keyposes[0].translation().y(), 0.19, kEpsilon); +} + +TEST(ComputeTopDownKeyposes, ReturnsAnEmptyPathForNoHeights) +{ + EXPECT_TRUE(computeTopDownKeyposes(makePose({ 0.5, 0.0, 0.115 }, 0.0), topDownGraspOrientation(0.0), + Eigen::Vector3d::Zero(), {}) + .empty()); +} + +TEST(ComputeTopDownKeyposes, TakesOnlyThePositionFromATiltedAimPose) +{ + // Cube poses come from live physics and are never exactly level. The waypoint must still sit + // straight above the cube, since the arm approaches vertically. + Eigen::Isometry3d tilted(Eigen::AngleAxisd(0.3, Eigen::Vector3d::UnitZ()) * + Eigen::AngleAxisd(0.05, Eigen::Vector3d::UnitX())); + tilted.translation() = Eigen::Vector3d(0.5, 0.0, 0.115); + + const auto keyposes = computeTopDownKeyposes(tilted, topDownGraspOrientation(0.3), Eigen::Vector3d::Zero(), { 0.12 }); + + ASSERT_EQ(keyposes.size(), 1u); + EXPECT_NEAR(keyposes[0].translation().x(), 0.5, kEpsilon); + EXPECT_NEAR(keyposes[0].translation().y(), 0.0, kEpsilon); + EXPECT_NEAR(keyposes[0].translation().z(), 0.235, kEpsilon); + const Eigen::Vector3d approach = keyposes[0].rotation() * Eigen::Vector3d::UnitZ(); + EXPECT_NEAR(approach.z(), -1.0, kEpsilon); +} + +int main(int argc, char** argv) +{ + testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/src/vla_sim_behaviors/test/test_plan_joint_spline_through_poses.cpp b/src/vla_sim_behaviors/test/test_plan_joint_spline_through_poses.cpp new file mode 100644 index 000000000..b939d5ced --- /dev/null +++ b/src/vla_sim_behaviors/test/test_plan_joint_spline_through_poses.cpp @@ -0,0 +1,181 @@ +// Copyright 2026 PickNik Inc. +// All rights reserved. +// +// Unauthorized copying of this code base via any medium is strictly prohibited. +// Proprietary and confidential. + +#include +#include + +#include + +#include + +namespace +{ +using vla_sim_behaviors::JointSpline; +using vla_sim_behaviors::segmentDuration; +using vla_sim_behaviors::splineDuration; +using vla_sim_behaviors::splineKnotParameters; + +Eigen::VectorXd vec(std::initializer_list values) +{ + Eigen::VectorXd result(static_cast(values.size())); + Eigen::Index i = 0; + for (const double value : values) + { + result[i++] = value; + } + return result; +} + +/** Two knots one unit apart on a single joint, over the full parameter span. */ +JointSpline twoKnotSpline() +{ + return JointSpline({ 0.0, 1.0 }, { vec({ 0.0 }), vec({ 1.0 }) }); +} + +/** Three knots, the interior one deliberately off the midpoint so corner-cutting shows. */ +JointSpline threeKnotSpline() +{ + return JointSpline({ 0.0, 0.25, 1.0 }, { vec({ 0.0, 0.0 }), vec({ 0.5, -1.0 }), vec({ 2.0, 1.0 }) }); +} +} // namespace + +TEST(JointSpline, RejectsFewerThanTwoKnots) +{ + EXPECT_THROW(JointSpline({ 0.0 }, { vec({ 0.0 }) }), std::invalid_argument); + EXPECT_THROW(JointSpline({}, {}), std::invalid_argument); +} + +TEST(JointSpline, RejectsParameterCountMismatch) +{ + EXPECT_THROW(JointSpline({ 0.0, 0.5, 1.0 }, { vec({ 0.0 }), vec({ 1.0 }) }), std::invalid_argument); +} + +TEST(JointSpline, RejectsNonIncreasingParameters) +{ + EXPECT_THROW(JointSpline({ 0.0, 0.0 }, { vec({ 0.0 }), vec({ 1.0 }) }), std::invalid_argument); + EXPECT_THROW(JointSpline({ 1.0, 0.0 }, { vec({ 0.0 }), vec({ 1.0 }) }), std::invalid_argument); +} + +TEST(JointSpline, RejectsKnotsOfDifferentWidths) +{ + EXPECT_THROW(JointSpline({ 0.0, 1.0 }, { vec({ 0.0 }), vec({ 1.0, 2.0 }) }), std::invalid_argument); +} + +TEST(JointSpline, PassesThroughEveryKnotExactly) +{ + const JointSpline spline = threeKnotSpline(); + EXPECT_NEAR(spline.position(0.0)[0], 0.0, 1e-12); + EXPECT_NEAR(spline.position(0.25)[0], 0.5, 1e-12); + EXPECT_NEAR(spline.position(0.25)[1], -1.0, 1e-12); + EXPECT_NEAR(spline.position(1.0)[0], 2.0, 1e-12); +} + +TEST(JointSpline, StartsAndEndsAtRest) +{ + const JointSpline spline = threeKnotSpline(); + EXPECT_NEAR(spline.velocity(0.0).cwiseAbs().maxCoeff(), 0.0, 1e-9); + EXPECT_NEAR(spline.velocity(1.0).cwiseAbs().maxCoeff(), 0.0, 1e-9); +} + +TEST(JointSpline, FlowsThroughInteriorKnotsWithoutStopping) +{ + // The reason for one spline over the whole chain rather than a clamped move per segment: + // chaining would pin the interior knot to zero velocity and stop the arm at every waypoint. + const JointSpline spline = threeKnotSpline(); + EXPECT_GT(spline.velocity(0.25).cwiseAbs().maxCoeff(), 0.5); +} + +TEST(JointSpline, IsContinuousInVelocityAcrossAnInteriorKnot) +{ + const JointSpline spline = threeKnotSpline(); + const Eigen::VectorXd before = spline.velocity(0.25 - 1e-7); + const Eigen::VectorXd after = spline.velocity(0.25 + 1e-7); + EXPECT_LT((after - before).cwiseAbs().maxCoeff(), 1e-5); +} + +TEST(JointSpline, PeaksAtOneAndAHalfTimesTheMeanOnASingleSegment) +{ + // Pins the profile: a clamped cubic peaks at 1.5x its mean speed. Smootherstep would be + // 1.875 and a linear ramp 1.0, so this fails if the interpolation is swapped out. + EXPECT_NEAR(twoKnotSpline().peakSpeed()[0], 1.5, 1e-9); + EXPECT_NEAR(twoKnotSpline().velocity(0.5)[0], 1.5, 1e-9); +} + +TEST(JointSpline, ReportsPeakSpeedPerJointIndependently) +{ + const JointSpline spline = JointSpline({ 0.0, 1.0 }, { vec({ 0.0, 0.0 }), vec({ 1.0, 4.0 }) }); + const Eigen::VectorXd peak = spline.peakSpeed(); + EXPECT_NEAR(peak[0], 1.5, 1e-9); + EXPECT_NEAR(peak[1], 6.0, 1e-9); +} + +TEST(JointSpline, SaturatesOutsideTheParameterRangeRatherThanExtrapolating) +{ + // A cubic run past its last knot diverges fast, so evaluating off the end holds the knot. + const JointSpline spline = twoKnotSpline(); + EXPECT_NEAR(spline.position(-0.5)[0], 0.0, 1e-12); + EXPECT_NEAR(spline.position(1.5)[0], 1.0, 1e-12); + EXPECT_NEAR(spline.velocity(1.5)[0], 0.0, 1e-9); +} + +TEST(SegmentDuration, TakesTheCartesianBudgetWhenTheJointCapIsSlack) +{ + // 0.2 m at 0.065 m/s is 3.08 s; the joint cap allows far quicker, so it does not bind. + EXPECT_NEAR(segmentDuration(0.2, vec({ 0.1, 0.1 }), 0.065, vec({ 0.695, 0.695 })), 0.2 / 0.065, 1e-9); +} + +TEST(SegmentDuration, TakesTheJointCapWhenThePathIsShortButTheArmTurnsFar) +{ + // A wrist flip in place: no Cartesian distance to pay for, but the joint still has to + // stay under its limit. + const double duration = segmentDuration(0.0, vec({ 3.0 }), 0.065, vec({ 0.695 })); + EXPECT_NEAR(duration, 1.875 * 3.0 / 0.695, 1e-9); +} + +TEST(SegmentDuration, IgnoresJointsWithoutAVelocityLimit) +{ + EXPECT_NEAR(segmentDuration(0.1, vec({ 3.0 }), 0.065, vec({ 0.0 })), 0.1 / 0.065, 1e-9); +} + +TEST(SplineKnotParameters, SpansZeroToOneInProportionToDuration) +{ + const std::vector parameters = splineKnotParameters({ 1.0, 3.0 }); + ASSERT_EQ(parameters.size(), 3u); + EXPECT_NEAR(parameters[0], 0.0, 1e-12); + EXPECT_NEAR(parameters[1], 0.25, 1e-12); + EXPECT_NEAR(parameters[2], 1.0, 1e-12); +} + +TEST(SplineKnotParameters, StaysStrictlyIncreasingThroughAZeroLengthSegment) +{ + // Two waypoints at the same place would otherwise collapse onto one parameter and make + // the spline unsolvable. + const std::vector parameters = splineKnotParameters({ 1.0, 0.0, 1.0 }); + EXPECT_GT(parameters[2], parameters[1]); + EXPECT_NO_THROW(JointSpline(parameters, { vec({ 0.0 }), vec({ 1.0 }), vec({ 1.0 }), vec({ 2.0 }) })); +} + +TEST(SplineDuration, SpendsTheCartesianLengthAtTheRequestedSpeed) +{ + // The oracle's regime: the tip speed budget binds and the joint caps stay slack, so the + // whole segment takes exactly length / speed. + const JointSpline spline = twoKnotSpline(); + EXPECT_NEAR(splineDuration(spline, 0.381, 0.065, vec({ 0.695 })), 0.381 / 0.065, 1e-9); +} + +TEST(SplineDuration, StretchesUntilThePeakJointSpeedFitsUnderTheCap) +{ + // Peak speed is 1.5 rad per unit parameter, so a 0.5 rad/s cap needs 3 s. + const JointSpline spline = twoKnotSpline(); + EXPECT_NEAR(splineDuration(spline, 0.0, 0.065, vec({ 0.5 })), 3.0, 1e-9); +} + +TEST(SplineDuration, ClampsToASaneRange) +{ + const JointSpline spline = twoKnotSpline(); + EXPECT_NEAR(splineDuration(spline, 1e-9, 0.065, vec({ 1e6 })), 0.2, 1e-9); + EXPECT_NEAR(splineDuration(spline, 1e6, 0.065, vec({ 1e6 })), 60.0, 1e-9); +} diff --git a/src/vla_sim_behaviors/test/test_send_gripper_command.cpp b/src/vla_sim_behaviors/test/test_send_gripper_command.cpp new file mode 100644 index 000000000..8da2ced3e --- /dev/null +++ b/src/vla_sim_behaviors/test/test_send_gripper_command.cpp @@ -0,0 +1,178 @@ +// Copyright 2026 PickNik Inc. +// All rights reserved. +// +// Unauthorized copying of this code base via any medium is strictly prohibited. +// Proprietary and confidential. + +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include + +#include + +namespace +{ +using vla_sim_behaviors::SendGripperCommand; +using GripperCommand = control_msgs::action::GripperCommand; + +constexpr auto kActionName = "/test_gripper_cmd"; +constexpr double kTargetPosition = 0.6; + +/** + * @brief A gripper server that accepts every goal and never finishes it. + * + * Closing onto an object behaves this way for seconds at a time: the jaws stall short of the + * commanded position and the goal stays live. A Behavior that waited would block for all of it. + */ +class StallingGripperServer +{ +public: + explicit StallingGripperServer(const rclcpp::Node::SharedPtr& node) + { + server_ = rclcpp_action::create_server( + node, kActionName, + [this](const rclcpp_action::GoalUUID&, std::shared_ptr goal) { + received_.set_value(goal->command.position); + return rclcpp_action::GoalResponse::ACCEPT_AND_EXECUTE; + }, + [](const std::shared_ptr>&) { + return rclcpp_action::CancelResponse::ACCEPT; + }, + [this](const std::shared_ptr>& handle) { + // Held, never completed, exactly like jaws still compressing an object. + std::lock_guard lock(mutex_); + held_ = handle; + }); + } + + /** @brief Ends the held goal, which rclcpp_action can wedge on if left executing. */ + ~StallingGripperServer() + { + // Terminating the goal re-enters the server's own callbacks, so the lock is + // released first; holding it here throws EDEADLK out of a destructor. + std::shared_ptr> held; + { + std::lock_guard lock(mutex_); + held = held_; + } + if (held != nullptr && held->is_executing()) + { + held->abort(std::make_shared()); + } + } + + /** @brief Commanded position of the first goal the server received. */ + [[nodiscard]] std::future receivedPosition() + { + return received_.get_future(); + } + +private: + rclcpp_action::Server::SharedPtr server_; + std::mutex mutex_; + std::shared_ptr> held_; + std::promise received_; +}; + +/** @brief Spins @p node on its own thread for the lifetime of the fixture. */ +class SendGripperCommandTest : public ::testing::Test +{ +public: + SendGripperCommandTest() + : node_(std::make_shared("send_gripper_command_test")) + , context_(std::make_shared(node_)) + { + executor_.add_node(node_); + spin_thread_ = std::thread([this] { executor_.spin(); }); + } + + ~SendGripperCommandTest() override + { + stopSpinning(); + } + + /** + * @brief Stops the executor and joins its thread. Idempotent. + * + * A test that owns an action server must call this before the server goes out of scope. + * Server callbacks run on the executor thread and touch the server's own members, so + * letting the server die first is a use-after-free: it aborts the whole binary on + * std::mutex::lock, taking the other tests with it. + */ + void stopSpinning() + { + executor_.cancel(); + if (spin_thread_.joinable()) + { + spin_thread_.join(); + } + } + + [[nodiscard]] static BT::NodeConfiguration configFor(const std::string& action_name, double timeout) + { + BT::NodeConfiguration config; + config.blackboard = BT::Blackboard::create(); + config.input_ports["gripper_command_action_name"] = action_name; + config.input_ports["position"] = std::to_string(kTargetPosition); + config.input_ports["max_effort"] = "0.0"; + config.input_ports["wait_for_server_timeout"] = std::to_string(timeout); + return config; + } + + rclcpp::Node::SharedPtr node_; + std::shared_ptr context_; + rclcpp::executors::MultiThreadedExecutor executor_; + std::thread spin_thread_; +}; +} // namespace + +TEST_F(SendGripperCommandTest, SucceedsWithoutWaitingForAGoalThatNeverFinishes) +{ + // GIVEN a gripper server that accepts the goal and never completes it + StallingGripperServer server(node_); + std::future received = server.receivedPosition(); + SendGripperCommand behavior("send_gripper", configFor(kActionName, 5.0), context_); + + // WHEN the Behavior is ticked + const BT::NodeStatus status = behavior.tick(); + + // THEN it reports success rather than blocking on the unfinished goal + EXPECT_EQ(status, BT::NodeStatus::SUCCESS); + + // AND the server received the commanded position + ASSERT_EQ(received.wait_for(std::chrono::seconds(10)), std::future_status::ready) + << "The goal never reached the action server."; + EXPECT_DOUBLE_EQ(received.get(), kTargetPosition); + + stopSpinning(); // before `server` is destroyed; see stopSpinning(). +} + +TEST_F(SendGripperCommandTest, FailsWhenNoActionServerIsAvailable) +{ + // GIVEN no server on the configured action name + SendGripperCommand behavior("send_gripper", configFor("/absent_gripper_cmd", 0.2), context_); + + // WHEN the Behavior is ticked + // THEN it fails rather than reporting a command nothing can execute + EXPECT_EQ(behavior.tick(), BT::NodeStatus::FAILURE); +} + +int main(int argc, char** argv) +{ + rclcpp::init(argc, argv); + testing::InitGoogleTest(&argc, argv); + const int result = RUN_ALL_TESTS(); + rclcpp::shutdown(); + return result; +} diff --git a/src/vla_sim_behaviors/test/test_wait_for_episode_start.cpp b/src/vla_sim_behaviors/test/test_wait_for_episode_start.cpp new file mode 100644 index 000000000..9ebf167cc --- /dev/null +++ b/src/vla_sim_behaviors/test/test_wait_for_episode_start.cpp @@ -0,0 +1,217 @@ +// Copyright 2026 PickNik Inc. +// All rights reserved. +// +// Unauthorized copying of this code base via any medium is strictly prohibited. +// Proprietary and confidential. + +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include + +#include + +namespace +{ +using vla_sim_behaviors::WaitForEpisodeStart; +using GetActiveRecording = moveit_studio_internal_msgs::srv::GetActiveRecording; + +constexpr auto kServiceName = "/test_active_recording"; + +/** @brief A Trainer stand-in whose reported session state the test controls. */ +class FakeTrainer +{ +public: + explicit FakeTrainer(const rclcpp::Node::SharedPtr& node) + { + server_ = node->create_service(kServiceName, + [this](const std::shared_ptr, + std::shared_ptr response) { + ++calls_; + std::lock_guard lock(mutex_); + response->status.success = true; + response->session_json = payload_; + }); + } + + void setState(const std::string& state) + { + std::lock_guard lock(mutex_); + payload_ = R"({"id":"rec-1","state":")" + state + R"(","current_episode":1})"; + } + + /** @brief Serve @p payload verbatim, for session_json a parser cannot use. */ + void setRawPayload(const std::string& payload) + { + std::lock_guard lock(mutex_); + payload_ = payload; + } + + [[nodiscard]] int calls() const + { + return calls_; + } + +private: + rclcpp::Service::SharedPtr server_; + // Written by the test thread, read by the executor's. + mutable std::mutex mutex_; + std::string payload_{ R"({"id":"rec-1","state":"preparing","current_episode":1})" }; + std::atomic_int calls_{ 0 }; +}; + +/** @brief Spins @p node on its own thread for the lifetime of the fixture. */ +class WaitForEpisodeStartTest : public ::testing::Test +{ +public: + WaitForEpisodeStartTest() + : node_(std::make_shared("wait_for_episode_start_test")) + , context_(std::make_shared(node_)) + { + executor_.add_node(node_); + spin_thread_ = std::thread([this] { executor_.spin(); }); + } + + ~WaitForEpisodeStartTest() override + { + executor_.cancel(); + if (spin_thread_.joinable()) + { + spin_thread_.join(); + } + } + + [[nodiscard]] static BT::NodeConfiguration configFor(const std::string& service_name, double timeout) + { + BT::NodeConfiguration config; + config.blackboard = BT::Blackboard::create(); + config.input_ports["service_name"] = service_name; + config.input_ports["timeout"] = std::to_string(timeout); + return config; + } + + /** @brief Ticks @p behavior until it leaves RUNNING, or returns RUNNING at @p deadline. + * + * executeTick() rather than tick(): the latter is StatefulActionNode's protected + * hook, and only executeTick() routes the first tick through onStart(). */ + static BT::NodeStatus tickUntilSettled(WaitForEpisodeStart& behavior, std::chrono::seconds deadline) + { + const auto give_up = std::chrono::steady_clock::now() + deadline; + BT::NodeStatus status = behavior.executeTick(); + while (status == BT::NodeStatus::RUNNING && std::chrono::steady_clock::now() < give_up) + { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + status = behavior.executeTick(); + } + return status; + } + + rclcpp::Node::SharedPtr node_; + std::shared_ptr context_; + rclcpp::executors::MultiThreadedExecutor executor_; + std::thread spin_thread_; +}; +} // namespace + +TEST_F(WaitForEpisodeStartTest, RunsUntilTheSessionReportsRecording) +{ + // GIVEN a session that is still preparing + FakeTrainer trainer(node_); + WaitForEpisodeStart behavior("wait_for_episode", configFor(kServiceName, 10.0), context_); + + // WHEN it is ticked while the session has not opened an episode + ASSERT_EQ(behavior.executeTick(), BT::NodeStatus::RUNNING); + std::this_thread::sleep_for(std::chrono::milliseconds(300)); + + // THEN it stays RUNNING rather than releasing the motion early + EXPECT_EQ(behavior.executeTick(), BT::NodeStatus::RUNNING); + + // WHEN the episode opens + trainer.setState("recording"); + + // THEN it succeeds + EXPECT_EQ(tickUntilSettled(behavior, std::chrono::seconds(10)), BT::NodeStatus::SUCCESS); + EXPECT_GT(trainer.calls(), 1) << "The state was read once and never re-polled."; +} + +TEST_F(WaitForEpisodeStartTest, FailsWhenNoEpisodeOpensWithinTheTimeout) +{ + // GIVEN a session that never leaves preparing + FakeTrainer trainer(node_); + WaitForEpisodeStart behavior("wait_for_episode", configFor(kServiceName, 0.5), context_); + + // WHEN it is ticked + // THEN it fails rather than waiting on a recording that will never start + EXPECT_EQ(tickUntilSettled(behavior, std::chrono::seconds(20)), BT::NodeStatus::FAILURE); + + // AND says which state it was stuck in, so the operator can tell a slow recorder + // from a session that never started + EXPECT_NE(context_->logger->consumeErrorLogBuffer().find("'preparing'"), std::string::npos); +} + +TEST_F(WaitForEpisodeStartTest, KeepsWaitingWhenTheSessionCannotBeRead) +{ + // GIVEN a Trainer answering with a payload no parser can use + FakeTrainer trainer(node_); + trainer.setRawPayload("{not json"); + WaitForEpisodeStart behavior("wait_for_episode", configFor(kServiceName, 10.0), context_); + + // WHEN it is ticked + ASSERT_EQ(behavior.executeTick(), BT::NodeStatus::RUNNING); + std::this_thread::sleep_for(std::chrono::milliseconds(300)); + + // THEN an unreadable payload is not mistaken for an open episode + EXPECT_EQ(behavior.executeTick(), BT::NodeStatus::RUNNING); + + // AND it recovers once the Trainer answers properly again + trainer.setState("recording"); + EXPECT_EQ(tickUntilSettled(behavior, std::chrono::seconds(10)), BT::NodeStatus::SUCCESS); +} + +TEST_F(WaitForEpisodeStartTest, HaltsWithoutWaitingOutTheTimeout) +{ + // GIVEN a long wait on a session that has not opened an episode + FakeTrainer trainer(node_); + WaitForEpisodeStart behavior("wait_for_episode", configFor(kServiceName, 600.0), context_); + ASSERT_EQ(behavior.executeTick(), BT::NodeStatus::RUNNING); + + // WHEN the tree halts it (haltNode(), since halt() is the protected hook) + const auto started = std::chrono::steady_clock::now(); + behavior.haltNode(); + const auto elapsed = std::chrono::steady_clock::now() - started; + + // THEN it returns promptly instead of blocking until the timeout expires + EXPECT_LT(elapsed, std::chrono::seconds(5)) + << "halt() waited " << std::chrono::duration_cast(elapsed).count() << "ms."; +} + +TEST_F(WaitForEpisodeStartTest, FailsWhenTheTrainerIsAbsent) +{ + // GIVEN no Trainer on the configured service name + WaitForEpisodeStart behavior("wait_for_episode", configFor("/absent_active_recording", 10.0), context_); + + // WHEN it is ticked + // THEN it fails instead of blocking on a service nothing serves + EXPECT_EQ(tickUntilSettled(behavior, std::chrono::seconds(30)), BT::NodeStatus::FAILURE); + + // AND names the service, since a typo here is the likely cause + EXPECT_NE(context_->logger->consumeErrorLogBuffer().find("/absent_active_recording"), std::string::npos); +} + +int main(int argc, char** argv) +{ + rclcpp::init(argc, argv); + testing::InitGoogleTest(&argc, argv); + const int result = RUN_ALL_TESTS(); + rclcpp::shutdown(); + return result; +} diff --git a/src/vla_sim_behaviors/vla_sim_behaviors_plugin_description.xml b/src/vla_sim_behaviors/vla_sim_behaviors_plugin_description.xml new file mode 100644 index 000000000..9eaec688b --- /dev/null +++ b/src/vla_sim_behaviors/vla_sim_behaviors_plugin_description.xml @@ -0,0 +1,7 @@ + + + +