diff --git a/.dockerignore b/.dockerignore index 5d4eed38d..485fb637c 100644 --- a/.dockerignore +++ b/.dockerignore @@ -2,3 +2,6 @@ build install log +# VLA checkpoints: keep multi-GB weights out of the workspace image build context. +src/vla_sim/models +src/vla_sim/hf_cache diff --git a/.gitignore b/.gitignore index f6ee3643d..aa0371f2b 100644 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,15 @@ log/ .vscode/ MUJOCO_LOG.TXT .ccache/ + +__pycache__/ + +# VLA checkpoint drop folder: multi-GB model weights must never be committed. +src/vla_sim/models/* +!src/vla_sim/models/.gitkeep +src/vla_sim/hf_cache/* +!src/vla_sim/hf_cache/.gitkeep + +# Machine-local launcher/workspace settings and secrets (HF_TOKEN etc.). +.env +.env.* diff --git a/README.md b/README.md index 35054b120..154a72cd2 100644 --- a/README.md +++ b/README.md @@ -30,6 +30,7 @@ git submodule foreach --recursive git lfs pull - `lab_sim` - `lunar_sim` - `phoebe_sim` +- `vla_sim` - `moveit_pro_franka_configs/franka_base_config` - `moveit_pro_kinova_configs/kinova_gen3_base_config` - `moveit_pro_kinova_configs/kinova_gen3_site_config` diff --git a/docker-compose.yaml b/docker-compose.yaml index d4ea095a4..190697fbb 100644 --- a/docker-compose.yaml +++ b/docker-compose.yaml @@ -9,8 +9,13 @@ services: # The base image that all MoveIt Pro services extend off of. Builds the user workspace. base: {} - # Starts the MoveIt Pro Agent and the Bridge between the Agent and the Web UI. - runtime: {} + # Starts the MoveIt Pro runtime: the Agent and its Bridge to the Web UI. + runtime: + environment: + # The vla_sim adapter's target; derived from the same + # 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 # Starts the robot drivers. drivers: @@ -26,3 +31,59 @@ services: volumes: # Allow access to host hardware e.g. RealSense cameras - /dev:/dev + + # LeRobot VLA inference server for the vla_sim config (src/vla_sim/docker/). + # The workspace half of the product's `inference_server` skeleton (see + # /opt/moveit_pro/docker-compose.yaml): the two same-named services merge, + # with the product side owning lifecycle, GPU runtime, security, and the + # published port. Only `moveit_pro run --with-inference-server` and + # `--only-inference-server` build or start it. Model selection lives in + # src/vla_sim/config/vla_serving.yaml, mounted at /vla_config; everything + # below is a host setting. + inference_server: + # Declared here too, so a launcher whose compose file carries no + # inference_server skeleton still gates this service instead of building + # and starting it on every plain run and build. + profiles: ["inference"] + build: + # Relative paths resolve against the compose project directory + # (/opt/moveit_pro when the launcher runs this), so workspace paths must + # go through MOVEIT_HOST_USER_WORKSPACE, same as the base service. + context: ${MOVEIT_HOST_USER_WORKSPACE:-.}/src/vla_sim/docker + dockerfile: Dockerfile.vla_inference_server + args: + TORCH_INDEX: ${VLA_TORCH_INDEX:-} + # Host facts only; model knobs live in vla_serving.yaml (mounted below). + environment: + # Token for gated/private Hugging Face downloads (the pi0.5 processor + # resolves the gated google/paligemma tokenizer even for local + # checkpoints; accept its license once with this token, or pre-seed the + # cache and set HF_HUB_OFFLINE=1). + - HF_TOKEN=${HF_TOKEN:-} + - HF_HUB_OFFLINE=${HF_HUB_OFFLINE:-0} + - HF_HOME=/hf + # Scratch home for the mapped user; only /hf persists. The image ships + # a passwd entry for uid 1000 only, so USER keeps getpass.getuser() + # (torch cache-dir setup) working when the host uid differs. + - HOME=/tmp + - USER=vla + volumes: + # Hub downloads (checkpoints, tokenizers) persist here across restarts. + # The default is a workspace-owned folder that exists in every clone, so + # docker never creates the bind source as root (a missing host path + # would be created root-owned and the non-root container could not + # write to it). Point VLA_HF_CACHE at an absolute path to reuse an + # existing Hugging Face cache instead. + - ${VLA_HF_CACHE:-${MOVEIT_HOST_USER_WORKSPACE:-.}/src/vla_sim/hf_cache}:/hf + # Local checkpoints: drop a checkpoint directory into + # src/vla_sim/models/ and set the vla_serving.yaml checkpoint to + # /models/. VLA_MODELS_DIR overrides the host folder + # for checkpoints stored elsewhere. + - ${VLA_MODELS_DIR:-${MOVEIT_HOST_USER_WORKSPACE:-.}/src/vla_sim/models}:/models:ro + # Model-serving config (checkpoint, fps, device, ...); edits picked up on + # restart, no image rebuild. + - ${MOVEIT_HOST_USER_WORKSPACE:-.}/src/vla_sim/config:/vla_config:ro + # The server script and its tests; edits take effect on restart, no + # image rebuild, and `docker exec ... python -m unittest` reaches the + # test suite (see src/vla_sim/docker/README.md). + - ${MOVEIT_HOST_USER_WORKSPACE:-.}/src/vla_sim/docker:/app:ro diff --git a/src/vla_sim/CMakeLists.txt b/src/vla_sim/CMakeLists.txt new file mode 100644 index 000000000..343ca4b10 --- /dev/null +++ b/src/vla_sim/CMakeLists.txt @@ -0,0 +1,39 @@ +cmake_minimum_required(VERSION 3.22) +project(vla_sim) + +find_package(ament_cmake REQUIRED) + +# A self-contained MoveIt Pro config: MuJoCo description (scene + arm/gripper), +# MoveIt/control config, objectives, waypoints, and the agent-bridge launch. +install( + DIRECTORY + config + description + launch + objectives + waypoints + DESTINATION + share/${PROJECT_NAME} +) + +install(PROGRAMS + script/get_action_chunk_adapter.py + DESTINATION lib/${PROJECT_NAME} +) + +if(BUILD_TESTING) + find_package(ament_lint_auto REQUIRED) + find_package(ament_cmake_pytest REQUIRED) + + # docker/test_vla_inference_server.py is not wired here: it needs + # torch/lerobot and runs inside the inference_server container instead + # (see docker/README.md). + ament_add_pytest_test(test_get_action_chunk_adapter + test/test_get_action_chunk_adapter.py + TIMEOUT 60 + ) + + ament_lint_auto_find_test_dependencies() +endif() + +ament_package() diff --git a/src/vla_sim/LICENSE b/src/vla_sim/LICENSE new file mode 100644 index 000000000..574ef0790 --- /dev/null +++ b/src/vla_sim/LICENSE @@ -0,0 +1,25 @@ +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 copyright holder 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. diff --git a/src/vla_sim/README.md b/src/vla_sim/README.md new file mode 100644 index 000000000..028607b56 --- /dev/null +++ b/src/vla_sim/README.md @@ -0,0 +1,18 @@ +# vla_sim + +A MoveIt Pro MuJoCo simulation of a Kinova Gen3 arm stacking colored cubes on +command, driven by a vision-language-action policy. The `Stack Cubes with the +VLA Policy` objective runs the policy, which is served over HTTP by the +`inference_server` container built from [`docker/`](docker/), where the setup +and serving instructions live. + +## Hardware requirements + +An NVIDIA GPU is recommended. When one is present, MoveIt Pro makes it +available to the inference server automatically. The stack still runs without +one, but inference moves to the CPU, where the default pi0.5 checkpoint might +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. + +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 new file mode 100644 index 000000000..4853da7a4 --- /dev/null +++ b/src/vla_sim/config/config.yaml @@ -0,0 +1,137 @@ +############################################################### +# +# This configures the robot to work with MoveIt Pro +# +############################################################### + +# Baseline hardware configuration parameters for MoveIt Pro. +# [Required] +hardware: + # If the MoveIt Pro Agent should launch the ros2 controller node. + # [Optional, default=True] + launch_control_node: True + + # If the MoveIt Pro Agent should launch the robot state publisher. + # This should be false if you are launching the robot state publisher as part of drivers. + # [Optional, default=True] + launch_robot_state_publisher: True + + # Parameters used to configure the robot description through XACRO. + # A URDF and SRDF are both required. + # [Required] + robot_description: + urdf: + package: "vla_sim" + path: "description/picknik_kinova_gen3.xacro" + srdf: + package: "vla_sim" + path: "config/moveit/picknik_kinova_gen3_base.srdf" + # Specify any additional parameters required for the URDF. + # Many of these are specific to the UR descriptions packages, and can be customized as needed. + # [Optional] + urdf_params: + - mujoco_model: "description/mujoco/cube_stack_scene.xml" + - mujoco_viewer: false + +# Sets ROS global params for launch. +# [Optional] +ros_global_params: + # Whether or not to use simulated time. + # [Optional, default=False] + use_sim_time: False + +# Configuration files for MoveIt. +# For more information, refer to https://moveit.picknik.ai/main/doc/how_to_guides/moveit_configuration/moveit_configuration_tutorial.html +# [Required] +moveit_params: + # Used by the Waypoint Manager to save joint states from this joint group. + joint_group_name: "manipulator" + + kinematics: + package: "vla_sim" + path: "config/moveit/pose_ik_distance.yaml" + sensors_3d: + package: "vla_sim" + path: "config/moveit/sensors_3d.yaml" + joint_limits: + package: "vla_sim" + path: "config/moveit/joint_limits.yaml" + pose_jog: + package: "vla_sim" + path: "config/moveit/pose_jog.yaml" + joint_jog: + package: "vla_sim" + path: "config/moveit/joint_jog.yaml" + + publish: + planning_scene: True + geometry_updates: True + state_updates: True + transforms_updates: True + + trajectory_execution: + manage_controllers: True + allowed_execution_duration_scaling: 2.0 + allowed_goal_duration_margin: 5.0 + allowed_start_tolerance: 0.01 + +# Configuration for launching ros2_control processes. +# [Required, if using ros2_control] +ros2_control: + config: + package: "vla_sim" + path: "config/control/picknik_kinova_gen3.ros2_control.yaml" + # MoveIt Pro will load and activate these controllers at start up to ensure they are available. + # If not specified, it is up to the user to ensure the appropriate controllers are active and available + # for running the application. + # [Optional, default=[]] + controllers_active_at_startup: + - "force_torque_sensor_broadcaster" + - "robotiq_gripper_controller" + - "joint_state_broadcaster" + # Load but do not start these controllers so they can be activated later if needed. + controllers_inactive_at_startup: + - "joint_trajectory_admittance_controller" + - "velocity_force_controller" + - "joint_velocity_controller" + # Any controllers here will not be spawned by MoveIt Pro. + # [Optional, default=[]] + controllers_not_managed: [] + # Optionally configure remapping rules to let multiple controllers receive commands on the same topic. + # [Optional, default=[]] + controller_shared_topics: [] + +# Configuration for loading behaviors and objectives. +# [Required] +objectives: + # List of plugins for loading custom behaviors. + # [Required] + behavior_loader_plugins: + # This plugin will load the core MoveIt Pro Behaviors. + # Add additional plugin loaders as needed. + core: + - "moveit_pro::behaviors::CoreBehaviorsLoader" + - "moveit_pro::behaviors::MTCCoreBehaviorsLoader" + - "moveit_pro::behaviors::VisionBehaviorsLoader" + - "moveit_pro::behaviors::ConverterBehaviorsLoader" + - "moveit_pro::behaviors::MujocoBehaviorsLoader" + # Specify source folder for objectives + # [Required] + objective_library_paths: + core_objectives: + package_name: "moveit_pro_objectives" + relative_path: "objectives/core" + motion_objectives: + package_name: "moveit_pro_objectives" + relative_path: "objectives/motion" + mujoco_objectives: + package_name: "moveit_pro_objectives" + relative_path: "objectives/mujoco" + sim_objectives: + package_name: "vla_sim" + relative_path: "objectives" + # Specify the location of the saved waypoints file. + # [Required] + waypoints_file: + package_name: "vla_sim" + relative_path: "waypoints/waypoints.yaml" diff --git a/src/vla_sim/config/control/picknik_kinova_gen3.ros2_control.yaml b/src/vla_sim/config/control/picknik_kinova_gen3.ros2_control.yaml new file mode 100644 index 000000000..c30cf6873 --- /dev/null +++ b/src/vla_sim/config/control/picknik_kinova_gen3.ros2_control.yaml @@ -0,0 +1,135 @@ +controller_manager: + ros__parameters: + update_rate: 1000 # Hz + joint_state_broadcaster: + type: joint_state_broadcaster/JointStateBroadcaster + joint_trajectory_admittance_controller: + type: joint_trajectory_admittance_controller/JointTrajectoryAdmittanceController + twist_controller: + type: picknik_twist_controller/PicknikTwistController + # The gripper controller just relays position commands + robotiq_gripper_controller: + type: position_controllers/GripperActionController + fault_controller: + type: picknik_reset_fault_controller/PicknikResetFaultController + force_torque_sensor_broadcaster: + type: force_torque_sensor_broadcaster/ForceTorqueSensorBroadcaster + velocity_force_controller: + type: velocity_force_controller/VelocityForceController + joint_velocity_controller: + type: joint_velocity_controller/JointVelocityController + +twist_controller: + ros__parameters: + joint: tcp + interface_names: + - twist.linear.x + - twist.linear.y + - twist.linear.z + - twist.angular.x + - twist.angular.y + - twist.angular.z + +robotiq_gripper_controller: + ros__parameters: + default: true + joint: robotiq_85_left_knuckle_joint + allow_stalling: true + stall_timeout: 0.05 + goal_tolerance: 0.02 + +force_torque_sensor_broadcaster: + ros__parameters: + sensor_name: fts + state_interface_names: + - force.x + - force.y + - force.z + - torque.x + - torque.y + - torque.z + frame_id: fts_link + + +joint_trajectory_admittance_controller: + ros__parameters: + planning_group_name: manipulator + sensor_frame: fts_link + ee_frame: grasp_link + ft_sensor_name: fts + default_path_tolerance: 0.5 + stop_accelerations: [3.0, 3.0, 3.0, 3.0, 5.0, 5.0, 5.0] + + +velocity_force_controller: + ros__parameters: + planning_group_name: manipulator + sensor_frame: fts_link + ee_frame: grasp_link + ft_sensor_name: fts + ft_force_deadband: 2.0 + ft_torque_deadband: 1.0 + # Conservative per-joint jog caps, all under the hardware limits in + # moveit/joint_limits.yaml. The wrist cap matches joint_velocity_controller's, + # so a pose jog near a wrist singularity cannot spike the small actuators + # to their ceiling. + max_joint_velocity: + - 0.524 + - 0.524 + - 0.524 + - 0.524 + - 0.6 + - 0.6 + - 0.6 + max_joint_acceleration: + - 52.4 + - 52.4 + - 52.4 + - 52.4 + - 52.4 + - 52.4 + - 52.4 + max_cartesian_velocity: + - 0.25 + - 0.25 + - 0.25 + - 1.5707 + - 1.5707 + - 1.5707 + max_cartesian_acceleration: + - 20.0 + - 20.0 + - 20.0 + - 40.0 + - 40.0 + - 40.0 + +joint_velocity_controller: + ros__parameters: + # Joint group to control. + planning_group_name: manipulator + # Maximum joint-space velocities. + max_joint_velocity: + - 0.7 + - 0.7 + - 0.7 + - 0.7 + - 0.6 + - 0.6 + - 0.6 + # Maximum joint-space accelerations. + max_joint_acceleration: + - 2.6 + - 2.6 + - 2.6 + - 2.6 + - 5.0 + - 5.0 + - 5.0 + command_interfaces: ["position"] + # Padding (in radians) to add to joint position limits as a safety margin. + joint_limit_position_tolerance: 0.02 + # Timeout in seconds after which the controller will stop motion if no new commands are received. + command_timeout: 0.2 + # Rate in Hz at which the controller will publish the state. + state_publish_rate: 20 diff --git a/src/vla_sim/config/moveit/joint_jog.yaml b/src/vla_sim/config/moveit/joint_jog.yaml new file mode 100644 index 000000000..6d80c7e9c --- /dev/null +++ b/src/vla_sim/config/moveit/joint_jog.yaml @@ -0,0 +1,4 @@ +# Planning groups to use in JointJog, and their corresponding JVC controllers. +# The number of elements in `planning_groups` and `controllers` must match. +planning_groups: ['manipulator'] +controllers: ['joint_velocity_controller'] diff --git a/src/vla_sim/config/moveit/joint_limits.yaml b/src/vla_sim/config/moveit/joint_limits.yaml new file mode 100644 index 000000000..a8703f8ae --- /dev/null +++ b/src/vla_sim/config/moveit/joint_limits.yaml @@ -0,0 +1,12 @@ +# Overrides the URDF dynamics limits for MoveIt planning and the ExecutePolicy +# velocity/acceleration gate. Values are the Kinova Gen3 hardware limits: 1.3963 rad/s +# for the large actuators (joints 1-4), 1.2218 rad/s for the small actuators (joints 5-7), +# 52.4 rad/s^2 acceleration. +joint_limits: + joint_1: {has_velocity_limits: true, max_velocity: 1.3963, has_acceleration_limits: true, max_acceleration: 52.4} + joint_2: {has_velocity_limits: true, max_velocity: 1.3963, has_acceleration_limits: true, max_acceleration: 52.4} + joint_3: {has_velocity_limits: true, max_velocity: 1.3963, has_acceleration_limits: true, max_acceleration: 52.4} + joint_4: {has_velocity_limits: true, max_velocity: 1.3963, has_acceleration_limits: true, max_acceleration: 52.4} + joint_5: {has_velocity_limits: true, max_velocity: 1.2218, has_acceleration_limits: true, max_acceleration: 52.4} + joint_6: {has_velocity_limits: true, max_velocity: 1.2218, has_acceleration_limits: true, max_acceleration: 52.4} + joint_7: {has_velocity_limits: true, max_velocity: 1.2218, has_acceleration_limits: true, max_acceleration: 52.4} diff --git a/src/vla_sim/config/moveit/picknik_kinova_gen3_base.srdf b/src/vla_sim/config/moveit/picknik_kinova_gen3_base.srdf new file mode 100644 index 000000000..3b9f48c16 --- /dev/null +++ b/src/vla_sim/config/moveit/picknik_kinova_gen3_base.srdf @@ -0,0 +1,172 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/src/vla_sim/config/moveit/pose_ik_distance.yaml b/src/vla_sim/config/moveit/pose_ik_distance.yaml new file mode 100644 index 000000000..69e4c7c6f --- /dev/null +++ b/src/vla_sim/config/moveit/pose_ik_distance.yaml @@ -0,0 +1,5 @@ +manipulator: + kinematics_solver: pose_ik_plugin/PoseIKPlugin + target_tolerance: 0.001 + solve_mode: "optimize_distance" + optimization_timeout: 0.005 diff --git a/src/vla_sim/config/moveit/pose_jog.yaml b/src/vla_sim/config/moveit/pose_jog.yaml new file mode 100644 index 000000000..5ba7f244a --- /dev/null +++ b/src/vla_sim/config/moveit/pose_jog.yaml @@ -0,0 +1,4 @@ +# Planning groups to use in PoseJog, and their corresponding VFC controllers. +# The number of elements in `planning_groups` and `controllers` must match. +planning_groups: ['manipulator'] +controllers: ['velocity_force_controller'] diff --git a/src/vla_sim/config/moveit/sensors_3d.yaml b/src/vla_sim/config/moveit/sensors_3d.yaml new file mode 100644 index 000000000..048625fb5 --- /dev/null +++ b/src/vla_sim/config/moveit/sensors_3d.yaml @@ -0,0 +1,24 @@ +sensors: + - scene_scan_camera +scene_scan_camera: + # The name of the Octomap updater plugin that we are using. + sensor_plugin: "moveit_studio_plugins/PointCloudServiceOctomapUpdater" + # Topic to use for the point cloud updater service + point_cloud_service_name: "/point_cloud_service" + # Points further than this will not be used (in meters). + max_range: 1.2 + # Choose one of every 'point_subsample' points (select all if set to 1). + point_subsample: 1 + # Should always be >= 1.0. Scale up collision shapes in the scene before excluding them from the octomap. + padding_scale: 1.0 + # Absolute padding around scaled collision shapes when excluding them from the octomap (in meters). + padding_offset: 0.05 + # The octomap representation will be updated at rate less than or equal to this value. + max_update_rate: 0.1 + +# Specifies the resolution at which the octomap is maintained (in meters). +octomap_resolution: 0.03 +# Specifies the coordinate frame in which the Octomap representation will be stored. +# Note! When an OccupancyMonitor instance is initialized by the PlanningSceneMonitor, +# this frame parameter will not be used. Instead, the frame defaults to the planning frame. +octomap_frame: "base_link" diff --git a/src/vla_sim/config/vla_serving.yaml b/src/vla_sim/config/vla_serving.yaml new file mode 100644 index 000000000..e4b750ab7 --- /dev/null +++ b/src/vla_sim/config/vla_serving.yaml @@ -0,0 +1,42 @@ +# Model-serving configuration for the vla_sim inference server, and the place to +# swap or re-tune the policy. Edits take effect when the inference_server +# container restarts; there is no image rebuild. +# +# Each knob resolves as: a vla_inference_server.py CLI flag, then this file, +# then a built-in default. A missing or empty file is fine; a malformed one +# parks the server in the error state. + +# LeRobot checkpoint to serve: a Hugging Face repo id (org/name), or a local +# directory placed in src/vla_sim/models/, which the server reads as +# /models/. +# This chooses the robot's actions, so point it only at a source you trust. +checkpoint: PickNikRobotics/pi05_kinova_gen3_cube_stack_sim + +# LeRobot policy family (pi05 | smolvla | ...). Empty reads the checkpoint's +# config.json "type"; set it only when that field is missing. +policy_class: "" + +# Training frame rate. The chunk plays at 1/fps seconds per step, so a wrong +# value rescales every commanded joint velocity. Keep the objective's dt at +# 1/fps. 0 reads the rate from train_config.json, which this checkpoint omits. +fps: 10.0 + +# Torch device: auto | cpu | cuda. auto takes a usable GPU and otherwise cpu. +# Asking for cuda on a host without one fails at startup rather than running an +# order of magnitude slower on cpu. +device: auto + +# Trained observation.state width: 7 arm joints plus 1 gripper. 0 trusts +# config.json, which for this checkpoint reports its padded 32-dim +# architecture width instead of the real one. +state_dim: 8 + +# Real-time-chunking soft guidance, in steps past the frozen prefix. Only used +# for requests that leave guidance_horizon at 0; the shipped objective sets its +# own on the ExecutePolicy port, so tune it there. +guidance_horizon: 8 + +# Real-time-chunking guidance-weight schedule, a lerobot RTCAttentionSchedule +# name. An invalid value parks the server in the error state naming the valid +# ones. +rtc_schedule: EXP diff --git a/src/vla_sim/description/mujoco/assets/kinova/base_link.STL b/src/vla_sim/description/mujoco/assets/kinova/base_link.STL new file mode 100644 index 000000000..5f6aff0d1 --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/kinova/base_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3e38517e653d58f4160cfdcee34c9dea983d8b3ecb8f2c147444f58235335260 +size 183884 diff --git a/src/vla_sim/description/mujoco/assets/kinova/base_link.mtl b/src/vla_sim/description/mujoco/assets/kinova/base_link.mtl new file mode 100644 index 000000000..a771bd2d5 --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/kinova/base_link.mtl @@ -0,0 +1,22 @@ +# Blender 4.1.1 MTL File: 'None' +# www.blender.org + +newmtl Blue_002 +Ns 250.000000 +Ka 1.000000 1.000000 1.000000 +Kd 0.000000 0.000000 1.000000 +Ks 0.500000 0.500000 0.500000 +Ke 0.000000 0.000000 0.000000 +Ni 1.000000 +d 1.000000 +illum 2 + +newmtl White_002 +Ns 250.000000 +Ka 1.000000 1.000000 1.000000 +Kd 1.000000 1.000000 1.000000 +Ks 0.500000 0.500000 0.500000 +Ke 0.000000 0.000000 0.000000 +Ni 1.000000 +d 1.000000 +illum 2 diff --git a/src/vla_sim/description/mujoco/assets/kinova/base_link.obj b/src/vla_sim/description/mujoco/assets/kinova/base_link.obj new file mode 100644 index 000000000..502aa742f --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/kinova/base_link.obj @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d209049c7784b700dacd3ff6b92e5a5e653b175b81cd40cdc3ac08c2c4115a63 +size 221550 diff --git a/src/vla_sim/description/mujoco/assets/kinova/bracelet_no_vision_link.STL b/src/vla_sim/description/mujoco/assets/kinova/bracelet_no_vision_link.STL new file mode 100644 index 000000000..705d64f18 --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/kinova/bracelet_no_vision_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5b2ec7dafbf6a3c9a2cac717d18132a8a1efd9fcb8f2c40d9c12c8ece7323ec8 +size 1334384 diff --git a/src/vla_sim/description/mujoco/assets/kinova/bracelet_no_vision_link.mtl b/src/vla_sim/description/mujoco/assets/kinova/bracelet_no_vision_link.mtl new file mode 100644 index 000000000..66af0bf77 --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/kinova/bracelet_no_vision_link.mtl @@ -0,0 +1,12 @@ +# Blender 4.1.1 MTL File: 'None' +# www.blender.org + +newmtl White +Ns 250.000000 +Ka 1.000000 1.000000 1.000000 +Kd 1.000000 1.000000 1.000000 +Ks 0.500000 0.500000 0.500000 +Ke 0.000000 0.000000 0.000000 +Ni 1.000000 +d 1.000000 +illum 2 diff --git a/src/vla_sim/description/mujoco/assets/kinova/bracelet_no_vision_link.obj b/src/vla_sim/description/mujoco/assets/kinova/bracelet_no_vision_link.obj new file mode 100644 index 000000000..a17b9d033 --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/kinova/bracelet_no_vision_link.obj @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:190b2a97e289f208cd1e9cc053d1fdd2d7bd147909e026784d3a854d473b0c20 +size 1362966 diff --git a/src/vla_sim/description/mujoco/assets/kinova/bracelet_with_vision_link.STL b/src/vla_sim/description/mujoco/assets/kinova/bracelet_with_vision_link.STL new file mode 100644 index 000000000..72f81c50b --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/kinova/bracelet_with_vision_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8bb05e420a88532b4eee3948a0047fb8e7f341c3e5cc971d54e97e2aaffcf878 +size 1164084 diff --git a/src/vla_sim/description/mujoco/assets/kinova/bracelet_with_vision_link.mtl b/src/vla_sim/description/mujoco/assets/kinova/bracelet_with_vision_link.mtl new file mode 100644 index 000000000..fdb35325d --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/kinova/bracelet_with_vision_link.mtl @@ -0,0 +1,12 @@ +# Blender 4.1.1 MTL File: 'None' +# www.blender.org + +newmtl white +Ns 250.000000 +Ka 1.000000 1.000000 1.000000 +Kd 1.000000 1.000000 1.000000 +Ks 0.500000 0.500000 0.500000 +Ke 0.000000 0.000000 0.000000 +Ni 1.000000 +d 1.000000 +illum 2 diff --git a/src/vla_sim/description/mujoco/assets/kinova/bracelet_with_vision_link.obj b/src/vla_sim/description/mujoco/assets/kinova/bracelet_with_vision_link.obj new file mode 100644 index 000000000..66914c558 --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/kinova/bracelet_with_vision_link.obj @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:479465e12fa09738aa1853d6033839e61bb2f8976088f2b180842feca3fe39d7 +size 1560243 diff --git a/src/vla_sim/description/mujoco/assets/kinova/end_effector_link.STL b/src/vla_sim/description/mujoco/assets/kinova/end_effector_link.STL new file mode 100644 index 000000000..3e2754219 --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/kinova/end_effector_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5b6fb58e61fa475939767d68a446f97f1bff02c0e5935a3ea8bb51e6515783d8 +size 80 diff --git a/src/vla_sim/description/mujoco/assets/kinova/forearm_link.STL b/src/vla_sim/description/mujoco/assets/kinova/forearm_link.STL new file mode 100644 index 000000000..0c616b5f7 --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/kinova/forearm_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a00d96d2c7ec37e9fe4c5fb6d3fa2b4386dfd48fea5ac25632f4cb83b3505105 +size 456284 diff --git a/src/vla_sim/description/mujoco/assets/kinova/forearm_link.mtl b/src/vla_sim/description/mujoco/assets/kinova/forearm_link.mtl new file mode 100644 index 000000000..b435e8789 --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/kinova/forearm_link.mtl @@ -0,0 +1,12 @@ +# Blender 4.1.1 MTL File: 'None' +# www.blender.org + +newmtl White.001 +Ns 250.000000 +Ka 1.000000 1.000000 1.000000 +Kd 1.000000 1.000000 1.000000 +Ks 0.500000 0.500000 0.500000 +Ke 0.000000 0.000000 0.000000 +Ni 1.000000 +d 1.000000 +illum 2 diff --git a/src/vla_sim/description/mujoco/assets/kinova/forearm_link.obj b/src/vla_sim/description/mujoco/assets/kinova/forearm_link.obj new file mode 100644 index 000000000..92be4a04e --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/kinova/forearm_link.obj @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:72520fe6116385b46ad20320e03f7b1ca8741ecab6b1cec196c554fdb903b58f +size 507752 diff --git a/src/vla_sim/description/mujoco/assets/kinova/half_arm_1_link.STL b/src/vla_sim/description/mujoco/assets/kinova/half_arm_1_link.STL new file mode 100644 index 000000000..00b636d56 --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/kinova/half_arm_1_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8df461a71780d8e2d09104a507584424279d85173adf14867e359f7e58b3bba7 +size 514184 diff --git a/src/vla_sim/description/mujoco/assets/kinova/half_arm_1_link.mtl b/src/vla_sim/description/mujoco/assets/kinova/half_arm_1_link.mtl new file mode 100644 index 000000000..ac03f8c26 --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/kinova/half_arm_1_link.mtl @@ -0,0 +1,12 @@ +# Blender 4.1.1 MTL File: 'None' +# www.blender.org + +newmtl White.002 +Ns 250.000000 +Ka 1.000000 1.000000 1.000000 +Kd 1.000000 1.000000 1.000000 +Ks 0.500000 0.500000 0.500000 +Ke 0.000000 0.000000 0.000000 +Ni 1.000000 +d 1.000000 +illum 2 diff --git a/src/vla_sim/description/mujoco/assets/kinova/half_arm_1_link.obj b/src/vla_sim/description/mujoco/assets/kinova/half_arm_1_link.obj new file mode 100644 index 000000000..e6d13e0d9 --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/kinova/half_arm_1_link.obj @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d5c370bf5b89c7aad4387a52c1feead7e70e0d00b515e2d90aa096bad930a7c0 +size 562125 diff --git a/src/vla_sim/description/mujoco/assets/kinova/half_arm_2_link.STL b/src/vla_sim/description/mujoco/assets/kinova/half_arm_2_link.STL new file mode 100644 index 000000000..375f68608 --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/kinova/half_arm_2_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8a9f38de21648804b3ab7076ec73726031f2ade7c95eb2541d070553e025407b +size 489184 diff --git a/src/vla_sim/description/mujoco/assets/kinova/half_arm_2_link.mtl b/src/vla_sim/description/mujoco/assets/kinova/half_arm_2_link.mtl new file mode 100644 index 000000000..07c787567 --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/kinova/half_arm_2_link.mtl @@ -0,0 +1,12 @@ +# Blender 4.1.1 MTL File: 'None' +# www.blender.org + +newmtl White.003 +Ns 250.000000 +Ka 1.000000 1.000000 1.000000 +Kd 1.000000 1.000000 1.000000 +Ks 0.500000 0.500000 0.500000 +Ke 0.000000 0.000000 0.000000 +Ni 1.000000 +d 1.000000 +illum 2 diff --git a/src/vla_sim/description/mujoco/assets/kinova/half_arm_2_link.obj b/src/vla_sim/description/mujoco/assets/kinova/half_arm_2_link.obj new file mode 100644 index 000000000..b8847e002 --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/kinova/half_arm_2_link.obj @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bff684ba381830badc7e51e402394e16648d05a16639fa9934dac566b33b6663 +size 485542 diff --git a/src/vla_sim/description/mujoco/assets/kinova/shoulder_link.STL b/src/vla_sim/description/mujoco/assets/kinova/shoulder_link.STL new file mode 100644 index 000000000..71dbb2133 --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/kinova/shoulder_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a98e232e5197f09e9cd520e7a450c21bcec74bc49e2a19ce7cf86c5dc25f2e18 +size 476984 diff --git a/src/vla_sim/description/mujoco/assets/kinova/shoulder_link.mtl b/src/vla_sim/description/mujoco/assets/kinova/shoulder_link.mtl new file mode 100644 index 000000000..37d1f301d --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/kinova/shoulder_link.mtl @@ -0,0 +1,12 @@ +# Blender 4.1.1 MTL File: 'None' +# www.blender.org + +newmtl White.004 +Ns 250.000000 +Ka 1.000000 1.000000 1.000000 +Kd 1.000000 1.000000 1.000000 +Ks 0.500000 0.500000 0.500000 +Ke 0.000000 0.000000 0.000000 +Ni 1.000000 +d 1.000000 +illum 2 diff --git a/src/vla_sim/description/mujoco/assets/kinova/shoulder_link.obj b/src/vla_sim/description/mujoco/assets/kinova/shoulder_link.obj new file mode 100644 index 000000000..beb8cdb93 --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/kinova/shoulder_link.obj @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4d304cba1afc3fc7266159a5e34c26385d30a1b64764d3f04100d1192698bf4f +size 424456 diff --git a/src/vla_sim/description/mujoco/assets/kinova/spherical_wrist_1_link.STL b/src/vla_sim/description/mujoco/assets/kinova/spherical_wrist_1_link.STL new file mode 100644 index 000000000..cc0413db9 --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/kinova/spherical_wrist_1_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0b97a5e3b2de535344f8e8d4176d77f0cdb3455b235f35430fa0413208856e66 +size 546484 diff --git a/src/vla_sim/description/mujoco/assets/kinova/spherical_wrist_1_link.mtl b/src/vla_sim/description/mujoco/assets/kinova/spherical_wrist_1_link.mtl new file mode 100644 index 000000000..8b4a00c2c --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/kinova/spherical_wrist_1_link.mtl @@ -0,0 +1,12 @@ +# Blender 4.1.1 MTL File: 'None' +# www.blender.org + +newmtl White.005 +Ns 250.000000 +Ka 1.000000 1.000000 1.000000 +Kd 1.000000 1.000000 1.000000 +Ks 0.500000 0.500000 0.500000 +Ke 0.000000 0.000000 0.000000 +Ni 1.000000 +d 1.000000 +illum 2 diff --git a/src/vla_sim/description/mujoco/assets/kinova/spherical_wrist_1_link.obj b/src/vla_sim/description/mujoco/assets/kinova/spherical_wrist_1_link.obj new file mode 100644 index 000000000..c92ff7ba8 --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/kinova/spherical_wrist_1_link.obj @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:322c381b28eefca630d170cd434182434c063b6a36ebc1871152ccb490fa1864 +size 585957 diff --git a/src/vla_sim/description/mujoco/assets/kinova/spherical_wrist_2_link.STL b/src/vla_sim/description/mujoco/assets/kinova/spherical_wrist_2_link.STL new file mode 100644 index 000000000..aa941bfe5 --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/kinova/spherical_wrist_2_link.STL @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:12271fb7d0ee78fb5642424cf5078be20f0476061e74203db964b674031ce879 +size 516984 diff --git a/src/vla_sim/description/mujoco/assets/kinova/spherical_wrist_2_link.mtl b/src/vla_sim/description/mujoco/assets/kinova/spherical_wrist_2_link.mtl new file mode 100644 index 000000000..d012aae9e --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/kinova/spherical_wrist_2_link.mtl @@ -0,0 +1,12 @@ +# Blender 4.1.1 MTL File: 'None' +# www.blender.org + +newmtl White.006 +Ns 250.000000 +Ka 1.000000 1.000000 1.000000 +Kd 1.000000 1.000000 1.000000 +Ks 0.500000 0.500000 0.500000 +Ke 0.000000 0.000000 0.000000 +Ni 1.000000 +d 1.000000 +illum 2 diff --git a/src/vla_sim/description/mujoco/assets/kinova/spherical_wrist_2_link.obj b/src/vla_sim/description/mujoco/assets/kinova/spherical_wrist_2_link.obj new file mode 100644 index 000000000..a38660556 --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/kinova/spherical_wrist_2_link.obj @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:873e40ba691e3b15dcbd65ab9aad604c27b5747bbb674b1f38bfe3e0b2f648e1 +size 532462 diff --git a/src/vla_sim/description/mujoco/assets/rafti_finger.mtl b/src/vla_sim/description/mujoco/assets/rafti_finger.mtl new file mode 100644 index 000000000..65a2fec44 --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/rafti_finger.mtl @@ -0,0 +1,12 @@ +# Blender 4.1.1 MTL File: 'part.blend' +# www.blender.org + +newmtl Material +Ns 250.000000 +Ka 1.000000 1.000000 1.000000 +Kd 0.034536 0.034536 0.034536 +Ks 0.500000 0.500000 0.500000 +Ke 0.000000 0.000000 0.000000 +Ni 1.450000 +d 1.000000 +illum 2 diff --git a/src/vla_sim/description/mujoco/assets/rafti_finger.obj b/src/vla_sim/description/mujoco/assets/rafti_finger.obj new file mode 100644 index 000000000..a1f64d567 --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/rafti_finger.obj @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e4f6548123631dacb30a0a54516c522cf9aa6c556d9fbc2d42bc4938742c2032 +size 8371 diff --git a/src/vla_sim/description/mujoco/assets/rafti_finger_collision_1.mtl b/src/vla_sim/description/mujoco/assets/rafti_finger_collision_1.mtl new file mode 100644 index 000000000..0fcf20253 --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/rafti_finger_collision_1.mtl @@ -0,0 +1,12 @@ +# Blender 4.1.1 MTL File: 'rafti_finger.blend' +# www.blender.org + +newmtl Material.001 +Ns 250.000000 +Ka 1.000000 1.000000 1.000000 +Kd 0.034536 0.034536 0.034536 +Ks 0.500000 0.500000 0.500000 +Ke 0.000000 0.000000 0.000000 +Ni 1.450000 +d 1.000000 +illum 2 diff --git a/src/vla_sim/description/mujoco/assets/rafti_finger_collision_1.obj b/src/vla_sim/description/mujoco/assets/rafti_finger_collision_1.obj new file mode 100644 index 000000000..fdf6a470c --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/rafti_finger_collision_1.obj @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7f1163f54bd7763f89c5a5ac2a2db71e6e7303d2676d60c3c34162234568e093 +size 682 diff --git a/src/vla_sim/description/mujoco/assets/rafti_finger_collision_2.mtl b/src/vla_sim/description/mujoco/assets/rafti_finger_collision_2.mtl new file mode 100644 index 000000000..0fcf20253 --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/rafti_finger_collision_2.mtl @@ -0,0 +1,12 @@ +# Blender 4.1.1 MTL File: 'rafti_finger.blend' +# www.blender.org + +newmtl Material.001 +Ns 250.000000 +Ka 1.000000 1.000000 1.000000 +Kd 0.034536 0.034536 0.034536 +Ks 0.500000 0.500000 0.500000 +Ke 0.000000 0.000000 0.000000 +Ni 1.450000 +d 1.000000 +illum 2 diff --git a/src/vla_sim/description/mujoco/assets/rafti_finger_collision_2.obj b/src/vla_sim/description/mujoco/assets/rafti_finger_collision_2.obj new file mode 100644 index 000000000..108f0754e --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/rafti_finger_collision_2.obj @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:086bc1bda883756dea3c9a5306da2016547a812adb22e70811de0956eb043a98 +size 1465 diff --git a/src/vla_sim/description/mujoco/assets/rafti_finger_collision_3.mtl b/src/vla_sim/description/mujoco/assets/rafti_finger_collision_3.mtl new file mode 100644 index 000000000..0fcf20253 --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/rafti_finger_collision_3.mtl @@ -0,0 +1,12 @@ +# Blender 4.1.1 MTL File: 'rafti_finger.blend' +# www.blender.org + +newmtl Material.001 +Ns 250.000000 +Ka 1.000000 1.000000 1.000000 +Kd 0.034536 0.034536 0.034536 +Ks 0.500000 0.500000 0.500000 +Ke 0.000000 0.000000 0.000000 +Ni 1.450000 +d 1.000000 +illum 2 diff --git a/src/vla_sim/description/mujoco/assets/rafti_finger_collision_3.obj b/src/vla_sim/description/mujoco/assets/rafti_finger_collision_3.obj new file mode 100644 index 000000000..6110aa1fd --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/rafti_finger_collision_3.obj @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4236029a257ad3f15ca5b32caa8cf8b40920dee83a1ae8ac9dadb4e8b006a40a +size 771 diff --git a/src/vla_sim/description/mujoco/assets/rafti_finger_collision_4.mtl b/src/vla_sim/description/mujoco/assets/rafti_finger_collision_4.mtl new file mode 100644 index 000000000..0fcf20253 --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/rafti_finger_collision_4.mtl @@ -0,0 +1,12 @@ +# Blender 4.1.1 MTL File: 'rafti_finger.blend' +# www.blender.org + +newmtl Material.001 +Ns 250.000000 +Ka 1.000000 1.000000 1.000000 +Kd 0.034536 0.034536 0.034536 +Ks 0.500000 0.500000 0.500000 +Ke 0.000000 0.000000 0.000000 +Ni 1.450000 +d 1.000000 +illum 2 diff --git a/src/vla_sim/description/mujoco/assets/rafti_finger_collision_4.obj b/src/vla_sim/description/mujoco/assets/rafti_finger_collision_4.obj new file mode 100644 index 000000000..a7dccfb5b --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/rafti_finger_collision_4.obj @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5708a8d40a1848936f725031f090109eca1632b61e01d04d683fd33ce96ffd95 +size 3890 diff --git a/src/vla_sim/description/mujoco/assets/robotiq_2f85/base.stl b/src/vla_sim/description/mujoco/assets/robotiq_2f85/base.stl new file mode 100644 index 000000000..22a9c2b3b --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/robotiq_2f85/base.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1c7b9f2bd92d705fc4e897c94e905973a3c05f406845e942229433deb7041453 +size 1712484 diff --git a/src/vla_sim/description/mujoco/assets/robotiq_2f85/finger_link.stl b/src/vla_sim/description/mujoco/assets/robotiq_2f85/finger_link.stl new file mode 100644 index 000000000..cbfd80ed7 --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/robotiq_2f85/finger_link.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d5ee95e62f8415bf5b6e503c831a958f5fc1990bf9b2865329ec38a28932727c +size 89084 diff --git a/src/vla_sim/description/mujoco/assets/robotiq_2f85/finger_tip_link.stl b/src/vla_sim/description/mujoco/assets/robotiq_2f85/finger_tip_link.stl new file mode 100644 index 000000000..9614b95e1 --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/robotiq_2f85/finger_tip_link.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:633793332c081641ce200df27a40643cc293b29956e3cb5cb29cb811c33ef1c7 +size 110484 diff --git a/src/vla_sim/description/mujoco/assets/robotiq_2f85/inner_knuckle_link.stl b/src/vla_sim/description/mujoco/assets/robotiq_2f85/inner_knuckle_link.stl new file mode 100644 index 000000000..01ec1e2a5 --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/robotiq_2f85/inner_knuckle_link.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4fb74a8a0d76c0e471cf19fd48bc676fd5b19123798e7a39eb6aa56869354283 +size 84884 diff --git a/src/vla_sim/description/mujoco/assets/robotiq_2f85/knuckle_link.stl b/src/vla_sim/description/mujoco/assets/robotiq_2f85/knuckle_link.stl new file mode 100644 index 000000000..0b3f90a5c --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/robotiq_2f85/knuckle_link.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bbd6e868b1778bead60d2516c7ede581d7ad4e431744b1828757d6ea5129c112 +size 67084 diff --git a/src/vla_sim/description/mujoco/assets/robotiq_2f85/pad.stl b/src/vla_sim/description/mujoco/assets/robotiq_2f85/pad.stl new file mode 100644 index 000000000..231413751 --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/robotiq_2f85/pad.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4af1fa8d9bb285abecbf5dfc88e74cb30f059a7a074c2030dc35e9ee4316019d +size 15084 diff --git a/src/vla_sim/description/mujoco/assets/robotiq_2f85/silicone_pad.stl b/src/vla_sim/description/mujoco/assets/robotiq_2f85/silicone_pad.stl new file mode 100644 index 000000000..6915244ae --- /dev/null +++ b/src/vla_sim/description/mujoco/assets/robotiq_2f85/silicone_pad.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ca20c0fa61e6d3ce7b04bed25360e048e31dfd88ace04187aa6d9e8f4ca3fd0f +size 15084 diff --git a/src/vla_sim/description/mujoco/cube_stack_scene.xml b/src/vla_sim/description/mujoco/cube_stack_scene.xml new file mode 100644 index 000000000..8ceea1d5b --- /dev/null +++ b/src/vla_sim/description/mujoco/cube_stack_scene.xml @@ -0,0 +1,202 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/src/vla_sim/description/mujoco/gen3_7dof.xml b/src/vla_sim/description/mujoco/gen3_7dof.xml new file mode 100644 index 000000000..dea5b8101 --- /dev/null +++ b/src/vla_sim/description/mujoco/gen3_7dof.xml @@ -0,0 +1,746 @@ + + + + diff --git a/src/vla_sim/description/picknik_kinova_gen3.xacro b/src/vla_sim/description/picknik_kinova_gen3.xacro new file mode 100644 index 000000000..467815320 --- /dev/null +++ b/src/vla_sim/description/picknik_kinova_gen3.xacro @@ -0,0 +1,682 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + 100000.0 + 100000.0 + + + + + 1e+5 + 1 + 0 + 0.2 + 0.002 + 0 + + + + + + + + + + + + + + + + + + + + + + + + 100000.0 + 100000.0 + + + + + 1e+5 + 1 + 0 + 0.2 + 0.002 + 0 + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + ${-2*pi} + ${2*pi} + + + 0.0 + + + + + + + -2.41 + 2.41 + + + 0.0 + + + + + + + ${-2*pi} + ${2*pi} + + + -3.14 + + + + + + + -2.66 + 2.66 + + + -2.51 + + + + + + + ${-2*pi} + ${2*pi} + + + 0.0 + + + + + + + -2.23 + 2.23 + + + 0.96 + + + + + + + ${-2*pi} + ${2*pi} + + + 1.57 + + + + + + + + 0.7929 + + + + + + + + + + + + + + + + + + + + + + + + + picknik_mujoco_ros/MujocoSystem + $(arg mujoco_model) + vla_sim + 10 + true + 4 + 60 + $(arg mujoco_viewer) + + + + diff --git a/src/vla_sim/docker/Dockerfile.vla_inference_server b/src/vla_sim/docker/Dockerfile.vla_inference_server new file mode 100644 index 000000000..f1ef72769 --- /dev/null +++ b/src/vla_sim/docker/Dockerfile.vla_inference_server @@ -0,0 +1,39 @@ +FROM python:3.12-slim + +# opencv needs libGL/glib at import time even headless-adjacent; install once here +# rather than debugging ImportErrors inside a slim image. +# hadolint ignore=DL3008 +RUN apt-get update && apt-get install -y --no-install-recommends \ + libgl1 libglib2.0-0 \ + && rm -rf /var/lib/apt/lists/* + +# TORCH_INDEX pre-installs torch from a specific wheel index before lerobot +# resolves it. Empty (default) lets lerobot pull the standard build, which is +# CUDA-enabled on Linux and also runs on CPU; set it to +# https://download.pytorch.org/whl/cpu for a much smaller CPU-only image on +# machines without an NVIDIA GPU (VLA_TORCH_INDEX in the workspace .env). +# torch/torchvision are pinned (within lerobot 0.6.0's >=2.7,<2.12 range) so +# builds are reproducible and the CPU-only pre-install cannot be re-resolved +# to a CUDA build by the lerobot install. +ARG TORCH_INDEX= +ARG TORCH_VERSION=2.11.0 +ARG TORCHVISION_VERSION=0.26.0 +RUN if [ -n "$TORCH_INDEX" ]; then \ + pip install --no-cache-dir --index-url "$TORCH_INDEX" \ + --extra-index-url https://pypi.org/simple \ + "torch==$TORCH_VERSION" "torchvision==$TORCHVISION_VERSION"; \ + fi \ + && pip install --no-cache-dir "lerobot[pi,smolvla]==0.6.0" \ + "torch==$TORCH_VERSION" "torchvision==$TORCHVISION_VERSION" + +COPY vla_inference_server.py /app/vla_inference_server.py +WORKDIR /app + +# Non-root default for a bare `docker run` (the compose service overrides the +# user anyway); a real passwd entry keeps getpass.getuser() working in torch's +# import-time cache-dir setup. +RUN useradd --create-home --uid 1000 vla +USER vla + +EXPOSE 8973 +ENTRYPOINT ["python", "vla_inference_server.py"] diff --git a/src/vla_sim/docker/README.md b/src/vla_sim/docker/README.md new file mode 100644 index 000000000..b4aa89b91 --- /dev/null +++ b/src/vla_sim/docker/README.md @@ -0,0 +1,99 @@ +# Inference server + +`vla_inference_server.py` serves a LeRobot checkpoint (pi0.5, SmolVLA, ...) +over HTTP for the `Stack Cubes with the VLA Policy` objective. The workspace +`docker-compose.yaml` completes MoveIt Pro's `inference_server` service with +this directory's image. Model and device selection live in +`../config/vla_serving.yaml`. + +## Running it + +The default checkpoint resolves the gated `google/paligemma` tokenizer on first +load, so export a token from an account that has accepted the +[PaliGemma license](https://huggingface.co/google/paligemma-3b-pt-224): + +```bash +export HF_TOKEN=hf_your_token_here +moveit_pro build +moveit_pro run -c vla_sim --with-inference-server +``` + +The first run builds the image and downloads the checkpoint into `../hf_cache/`; +later runs reuse both. Then run **Stack Cubes with the VLA Policy** in the web +UI, and **Reset MuJoCo Sim** between attempts. + +Model loading takes a minute or more. To keep the model warm across restarts of +the stack, run the server on its own in one terminal and the stack, without +`--with-inference-server`, in another: + +```bash +# Terminal 1: the server, which prints its loading and ready status. +moveit_pro run --only-inference-server +# Terminal 2: restart this as often as you like; the loaded model survives. +moveit_pro run -c vla_sim +``` + +Pick one mode per session. Passing `--with-inference-server` while a +side-started server is running adopts that container, so stopping the stack +stops the server too. + +Serving a different checkpoint also takes two edits in +`../objectives/stack_cubes_with_the_vla_policy.xml`, because the request has to +match what the checkpoint was trained on: set `image_names` to its camera names, +which the server rejects the request for if they differ, and set `dt` to 1/`fps`. + +## Environment + +Set these in the workspace `.env`; all are optional. + +| Variable | Effect | +| --- | --- | +| `HF_TOKEN` | Token for gated or private Hugging Face downloads. | +| `HF_HUB_OFFLINE` | `1` serves only what is already in the cache, with no network access. | +| `VLA_HF_CACHE` | Host path for the Hugging Face cache. Defaults to `../hf_cache`. | +| `VLA_MODELS_DIR` | Host folder mounted at `/models`, for checkpoints stored outside the workspace. Defaults to `../models`. | +| `VLA_TORCH_INDEX` | Package index the image installs torch from, for example `https://download.pytorch.org/whl/cpu` on a machine with no NVIDIA GPU. Defaults to PyPI. | + +## The HTTP contract + +`GET /health` reports `loading` / `ready` / `error` and needs no token. +`POST /infer` requires the deployment's `MOVEIT_FRONTEND_KEY` as a bearer +token. The server speaks plain HTTP and publishes on `127.0.0.1` only, which is +what keeps that token off the network. Two settings decide what code and weights +the container runs, so point both only at sources you trust: `checkpoint` +chooses the robot's actions, and `VLA_TORCH_INDEX` supplies the torch build. + +## Running the image outside compose + +Compose builds the image and supplies the environment it needs. By hand, from +this directory: + +```bash +docker build -f Dockerfile.vla_inference_server -t vla_inference_server . +docker run --rm --user "$(id -u):$(id -g)" \ + --gpus all \ + -e HOME=/tmp -e USER=vla \ + -v "$PWD/../hf_cache:/hf" -e HF_HOME=/hf -e HF_TOKEN="$HF_TOKEN" \ + -v "$PWD/../config:/vla_config:ro" \ + -e MOVEIT_FRONTEND_KEY=moveit-secret-key \ + -p 127.0.0.1:8973:8973 vla_inference_server +``` + +`--user` keeps bind-mounted files from being written as uid 1000, which means +the image's own passwd entry no longer applies, so `HOME` and `USER` have to be +set for torch's import-time cache setup. `--gpus all` exposes the GPU to the +container; without it `device: auto` silently serves on cpu. Omit the flag on a +machine without an NVIDIA GPU, where it fails outright. The Hugging Face cache +mount makes checkpoint downloads persist, and the config mount is where the +server reads which checkpoint to load. The image's entrypoint already runs the +server, so anything after the image name is appended as arguments to it, and +`--checkpoint ` overrides the config. + +## Tests + +`test_vla_inference_server.py` needs `lerobot` and `torch`, so it runs in the +container, which mounts this directory at `/app`: + +```bash +docker exec "$(docker ps -qf name=inference_server)" python -m unittest -v test_vla_inference_server +``` diff --git a/src/vla_sim/docker/test_vla_inference_server.py b/src/vla_sim/docker/test_vla_inference_server.py new file mode 100644 index 000000000..0a46d5948 --- /dev/null +++ b/src/vla_sim/docker/test_vla_inference_server.py @@ -0,0 +1,820 @@ +#!/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. + +"""Tests for vla_inference_server.py: resolvers, image decoding, and the HTTP state machine. + +Runs in the same Python environment as vla_inference_server.py itself (lerobot/torch/cv2), +not the ROS workspace's pytest suite (see README.md for how to run this). +""" + +import argparse +import base64 +import http.client +import json +import os +import tempfile +import threading +import unittest +from unittest.mock import patch +from http.server import ThreadingHTTPServer + +import cv2 +import numpy as np +import torch +import yaml +from lerobot.configs.types import RTCAttentionSchedule +from lerobot.processor import RenameObservationsProcessorStep + +from vla_inference_server import ( + REQUEST_SOCKET_TIMEOUT_SECONDS, + ServerState, + apply_frontend_key, + decode_image_b64, + hub_access_error_message, + load_policy, + load_serving_config, + make_handler, + native_camera_map, + parse_args, + request_camera_names, + resolve_default, + resolve_device, + resolve_fps, + resolve_rtc_horizon, + resolve_rtc_schedule, +) + + +def encode_bgr_jpeg_b64(bgr: np.ndarray) -> str: + ok, buf = cv2.imencode(".jpg", bgr) + assert ok + return base64.b64encode(buf.tobytes()).decode("ascii") + + +class TestResolveDevice(unittest.TestCase): + """resolve_device: auto-selection and fail-loud explicit requests.""" + + def test_auto_prefers_cuda_when_available(self) -> None: + """device=auto on a GPU host serves on cuda, never silently on cpu.""" + self.assertEqual(resolve_device("auto", cuda_available=True), "cuda") + + def test_auto_falls_back_to_cpu(self) -> None: + """device=auto without a GPU serves on cpu.""" + self.assertEqual(resolve_device("auto", cuda_available=False), "cpu") + + def test_explicit_cuda_without_gpu_raises(self) -> None: + """An explicit cuda request on a CPU-only host is a startup error.""" + with self.assertRaises(ValueError): + resolve_device("cuda", cuda_available=False) + + def test_explicit_cpu_always_honored(self) -> None: + """An explicit cpu request is honored even when a GPU exists.""" + self.assertEqual(resolve_device("cpu", cuda_available=True), "cpu") + + +class TestLoadServingConfig(unittest.TestCase): + """load_serving_config: tolerant of missing/empty, loud on malformed.""" + + def _write(self, text: str) -> str: + with tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False) as handle: + handle.write(text) + path = handle.name + self.addCleanup(os.unlink, path) + return path + + def test_missing_file_returns_empty(self) -> None: + """A path with no file yields {}, so env and built-in defaults still apply.""" + # GIVEN a path that does not exist + # WHEN loading the serving config + result = load_serving_config("/nonexistent/vla_serving.yaml") + + # THEN it is an empty dict, not an error + self.assertEqual(result, {}) + + def test_empty_file_returns_empty(self) -> None: + """An empty YAML file is tolerated the same as a missing one.""" + # GIVEN an empty file + path = self._write("") + + # WHEN loading it + # THEN it yields {} rather than raising + self.assertEqual(load_serving_config(path), {}) + + def test_valid_file_parses_each_knob_with_native_type(self) -> None: + """A well-formed file parses keys with their YAML-native types.""" + # GIVEN a well-formed serving config + path = self._write("checkpoint: org/model\nfps: 10.0\nstate_dim: 8\n") + + # WHEN loading it + config = load_serving_config(path) + + # THEN each knob carries its native type + self.assertEqual(config["checkpoint"], "org/model") + self.assertEqual(config["fps"], 10.0) + self.assertEqual(config["state_dim"], 8) + + def test_malformed_file_raises(self) -> None: + """Broken YAML fails loudly instead of silently serving the wrong model.""" + # GIVEN a syntactically broken YAML file + path = self._write("checkpoint: [unterminated\n") + + # WHEN loading it + # THEN it raises, so the loader thread can park in the error state + with self.assertRaises(yaml.YAMLError): + load_serving_config(path) + + def test_non_mapping_file_raises(self) -> None: + """A top-level list is rejected, since knobs are looked up by key.""" + # GIVEN a YAML file whose top level is a list + path = self._write("- checkpoint\n- fps\n") + + # WHEN loading it + # THEN it is rejected with a message naming the expected shape + with self.assertRaises(ValueError): + load_serving_config(path) + + +class TestResolveDefault(unittest.TestCase): + """resolve_default: YAML > built-in for an argparse default.""" + + def test_yaml_beats_builtin(self) -> None: + """A YAML value wins over the built-in.""" + # GIVEN a YAML value + # WHEN resolving the default + # THEN the YAML value is used + self.assertEqual(resolve_default("cpu", "auto"), "cpu") + + def test_builtin_used_when_yaml_absent(self) -> None: + """With no YAML value, the built-in is returned.""" + # GIVEN no YAML value + # WHEN resolving + # THEN the built-in default is returned + self.assertEqual(resolve_default(None, "auto"), "auto") + + def test_yaml_zero_is_honored_over_builtin(self) -> None: + """A YAML value of 0 (the fps/state_dim auto sentinel) is honored, not skipped.""" + # GIVEN a YAML value of 0 and a non-zero built-in + # WHEN resolving + # THEN 0 is returned, not treated as absent + self.assertEqual(resolve_default(0, 8), 0) + + +class TestParseArgsCoercion(unittest.TestCase): + """parse_args: type-invalid YAML values defer to the error state, never crash.""" + + def test_non_numeric_yaml_values_park_in_config_error(self) -> None: + """A non-numeric fps or state_dim in the YAML lands in config_error with the + built-in default applied, so the socket still binds and the loader thread + reports the typo through /health instead of a pre-bind crash loop.""" + # GIVEN a serving config whose fps and state_dim are not numbers + with tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False) as f: + f.write("fps: ten\nstate_dim: [7, 1]\n") + path = f.name + try: + # WHEN parsing arguments against that config + with patch("sys.argv", ["vla_inference_server.py", "--config", path]): + args = parse_args() + finally: + os.unlink(path) + + # THEN both bad values are reported and the built-ins are used + self.assertIn("fps", args.config_error) + self.assertIn("state_dim", args.config_error) + self.assertEqual(args.fps, 0.0) + self.assertEqual(args.state_dim, 0) + + +class TestLoadPolicyMissingCheckpoint(unittest.TestCase): + """load_policy: an unset checkpoint parks the error state, never exits.""" + + def test_empty_checkpoint_parks_error_state(self) -> None: + """With no checkpoint configured, the loader thread parks in the error + state naming the fix, so the socket stays bound and /health plus the + objective's UI messages report it instead of the process exiting.""" + # GIVEN parsed args with a readable config but no checkpoint + state = ServerState() + args = argparse.Namespace(config_error="", checkpoint="") + + # WHEN the loader runs + load_policy(state, args) + + # THEN the server is parked in the error state with actionable detail + self.assertEqual(state.status, "error") + self.assertIn("checkpoint", state.detail) + self.assertIn("vla_serving.yaml", state.detail) + + +class TestResolveFps(unittest.TestCase): + """resolve_fps: explicit value wins; unresolvable rate is an error.""" + + def test_explicit_fps_wins(self) -> None: + """A positive fps skips the checkpoint lookup entirely.""" + self.assertEqual(resolve_fps("/nonexistent", 10.0), 10.0) + + def test_unresolvable_fps_raises(self) -> None: + """fps=0 with no readable train_config.json is a startup error, not a default.""" + with self.assertRaises(ValueError): + resolve_fps("nonexistent-checkpoint", 0.0) + + +class TestHubAccessErrorMessage(unittest.TestCase): + """hub_access_error_message: each access failure names its own fix.""" + + def test_gated_without_token_says_export_it(self) -> None: + """A gated repo with no token points at exporting HF_TOKEN.""" + message = hub_access_error_message("org/model", gated=True, token_present=False) + self.assertIn("HF_TOKEN is not set", message) + self.assertIn("export HF_TOKEN", message) + + def test_gated_with_token_says_accept_the_license(self) -> None: + """A gated repo with a token present points at license acceptance, not the token.""" + message = hub_access_error_message("org/model", gated=True, token_present=True) + self.assertIn("has not been granted access", message) + self.assertNotIn("export HF_TOKEN", message) + + def test_not_found_without_token_mentions_private_repos(self) -> None: + """An unknown repo without a token flags both a typo and the private case.""" + message = hub_access_error_message( + "org/model", gated=False, token_present=False + ) + self.assertIn("vla_serving.yaml", message) + self.assertIn("private repo", message) + + def test_not_found_with_token_points_at_the_name(self) -> None: + """An unknown repo with a token present points at the checkpoint name.""" + message = hub_access_error_message("org/typo", gated=False, token_present=True) + self.assertIn("org/typo", message) + self.assertIn("vla_serving.yaml", message) + + +class TestResolveRtcHorizon(unittest.TestCase): + """resolve_rtc_horizon: the service's guidance width -> lerobot's absolute horizon.""" + + def test_request_width_extends_past_the_prefix(self) -> None: + """The guided region ends inference_delay + width steps into the chunk, so a + width smaller than the delay can never shrink the frozen prefix.""" + self.assertEqual(resolve_rtc_horizon(9, 8, 12), 17) + + def test_zero_width_uses_the_server_default(self) -> None: + """guidance_horizon=0 defers to the server's configured width, per the contract.""" + self.assertEqual(resolve_rtc_horizon(9, 0, 12), 21) + + def test_zero_delay_passes_the_width_through(self) -> None: + """With no frozen prefix the horizon is just the guidance width.""" + self.assertEqual(resolve_rtc_horizon(0, 8, 12), 8) + + +class TestResolveRtcSchedule(unittest.TestCase): + """resolve_rtc_schedule: named schedules resolve; typos name the valid values.""" + + def test_known_schedule_resolves(self) -> None: + """A valid schedule name maps onto lerobot's enum.""" + self.assertEqual(resolve_rtc_schedule("EXP"), RTCAttentionSchedule.EXP) + + def test_unknown_schedule_names_the_valid_values(self) -> None: + """A typo'd schedule fails with a message listing the valid names and + pointing at the knob's file.""" + with self.assertRaises(ValueError) as ctx: + resolve_rtc_schedule("exp") + self.assertIn("EXP", str(ctx.exception)) + self.assertIn("vla_serving.yaml", str(ctx.exception)) + + +class TestDecodeImageB64(unittest.TestCase): + """decode_image_b64: base64 JPEG -> CHW float32 [0,1] RGB tensor.""" + + def test_invalid_base64_raises(self) -> None: + """Malformed image bytes fail loudly (ValueError) instead of returning garbage.""" + with self.assertRaises(ValueError): + decode_image_b64(base64.b64encode(b"not a jpeg").decode("ascii")) + + def test_output_shape_and_dtype(self) -> None: + """A 4x2 BGR frame decodes to a (3, 4, 2) float32 tensor scaled to [0, 1].""" + bgr = np.zeros((4, 2, 3), dtype=np.uint8) + tensor = decode_image_b64(encode_bgr_jpeg_b64(bgr)) + + self.assertEqual(tuple(tensor.shape), (3, 4, 2)) + self.assertEqual(tensor.dtype, torch.float32) + self.assertGreaterEqual(float(tensor.min()), 0.0) + self.assertLessEqual(float(tensor.max()), 1.0) + + def test_bgr_to_rgb_channel_order(self) -> None: + """A pure-blue BGR frame decodes with the red channel near zero (BGR -> RGB swap).""" + bgr = np.zeros((8, 8, 3), dtype=np.uint8) + bgr[:, :, 0] = 255 # BGR channel 0 = blue + tensor = decode_image_b64(encode_bgr_jpeg_b64(bgr)) + + # channel 0 = red after the BGR->RGB swap, so it should stay dark despite the + # source being fully saturated on the blue channel; channel 2 = blue, saturated. + self.assertLess(float(tensor[0].mean()), 0.2) + self.assertGreater(float(tensor[2].mean()), 0.8) + + +class TestNativeCameraMap(unittest.TestCase): + """native_camera_map: dataset-native camera names from the preprocessor pipeline.""" + + def test_rename_step_yields_prefix_stripped_image_map(self) -> None: + """Image entries lose the feature prefix; non-image entries are ignored.""" + step = RenameObservationsProcessorStep( + rename_map={ + "observation.images.overview": "observation.images.base_0_rgb", + "observation.images.scene": "observation.images.right_wrist_0_rgb", + "observation.env_state": "observation.state", + } + ) + self.assertEqual( + native_camera_map([step]), + {"overview": "base_0_rgb", "scene": "right_wrist_0_rgb"}, + ) + + def test_pipeline_without_rename_step_yields_empty_map(self) -> None: + """A checkpoint whose dataset already used the slot names offers no aliases.""" + self.assertEqual(native_camera_map([object()]), {}) + + def test_first_image_renaming_step_wins_and_warns(self) -> None: + """With two image-renaming steps the first defines the request names, + and the ambiguity is logged next to the load's request-names line.""" + first = RenameObservationsProcessorStep( + rename_map={"observation.images.front": "observation.images.scene"} + ) + second = RenameObservationsProcessorStep( + rename_map={"observation.images.top": "observation.images.scene"} + ) + with patch("vla_inference_server.log") as mock_log: + result = native_camera_map([first, second]) + + self.assertEqual(result, {"front": "scene"}) + self.assertIn("WARNING", mock_log.call_args[0][0]) + + def test_non_image_rename_step_does_not_mask_a_later_image_one(self) -> None: + """A step renaming only state keys is skipped; the image-renaming step + behind it still defines the camera names, with no ambiguity warning.""" + state_only = RenameObservationsProcessorStep( + rename_map={"observation.env_state": "observation.state"} + ) + images = RenameObservationsProcessorStep( + rename_map={"observation.images.front": "observation.images.scene"} + ) + with patch("vla_inference_server.log") as mock_log: + result = native_camera_map([state_only, images]) + + self.assertEqual(result, {"front": "scene"}) + mock_log.assert_not_called() + + +class TestRequestCameraNames(unittest.TestCase): + """request_camera_names: the /infer image keys for a checkpoint, in order.""" + + def test_partial_rename_mixes_native_and_slot_names(self) -> None: + """A camera the checkpoint renames takes its dataset name; one it does + not rename keeps its slot name, in checkpoint-declared order.""" + self.assertEqual( + request_camera_names(["scene", "aux"], {"front": "scene"}), + ["front", "aux"], + ) + + def test_no_rename_map_keeps_slot_names(self) -> None: + """Without a rename step the config.json slot names are the request names.""" + self.assertEqual(request_camera_names(["a", "b"], {}), ["a", "b"]) + + def test_many_to_one_rename_warns_and_uses_the_last(self) -> None: + """Two dataset names mapping onto one slot cannot both be honored; the + collision is logged and the last one becomes the request name.""" + with patch("vla_inference_server.log") as mock_log: + result = request_camera_names(["scene"], {"a": "scene", "b": "scene"}) + + self.assertEqual(result, ["b"]) + self.assertIn("WARNING", mock_log.call_args[0][0]) + + +class FakeRunner: + """Stands in for PolicyRunner: same expected_state_dim/infer contract, no model.""" + + def __init__( + self, + infer_error: Exception | None = None, + camera_keys: list | None = None, + native_map: dict | None = None, + ) -> None: + self.device = "cpu" + self._infer_error = infer_error + # Like PolicyRunner, derived once at construction. + self.request_names = request_camera_names( + camera_keys if camera_keys is not None else ["scene"], + native_map if native_map is not None else {}, + ) + + def expected_state_dim(self) -> int: + return 2 + + def infer( + self, images, state, prompt, prev_chunk, inference_delay, guidance_horizon + ): + if self._infer_error is not None: + raise self._infer_error + return np.array([[0.1, 0.2]]), np.array([[0.5, 0.5]]) + + +class TestApplyFrontendKey(unittest.TestCase): + """Fail-closed handling of the MOVEIT_FRONTEND_KEY environment value.""" + + def test_blank_or_missing_key_parks_error_state(self) -> None: + """An unset or blank key parks the server so /health names the fix.""" + for raw_key in (None, "", " "): + state = ServerState() + self.assertFalse(apply_frontend_key(state, raw_key)) + self.assertEqual(state.status, "error") + self.assertIn("MOVEIT_FRONTEND_KEY", state.detail) + # /health serves the detail without a token; it must never name a + # usable key value, only point at the docs. + self.assertNotIn("moveit-secret-key", state.detail) + + def test_valid_key_is_stored_stripped(self) -> None: + """A usable key is stored without surrounding whitespace.""" + state = ServerState() + self.assertTrue(apply_frontend_key(state, " secret-key \n")) + self.assertEqual(state.frontend_key, "secret-key") + self.assertEqual(state.status, "loading") + + +class TestHttpStateMachine(unittest.TestCase): + """/health and /infer across the loading -> ready/error lifecycle.""" + + # Auth key served by every test server; _infer presents it by default. + TEST_KEY = "test-frontend-key" + + def _start(self, state: ServerState) -> http.client.HTTPConnection: + state.frontend_key = self.TEST_KEY + httpd = ThreadingHTTPServer(("127.0.0.1", 0), make_handler(state)) + threading.Thread(target=httpd.serve_forever, daemon=True).start() + # LIFO: shutdown() stops the serve loop first, then server_close() + # frees the listening socket. + self.addCleanup(httpd.server_close) + self.addCleanup(httpd.shutdown) + return http.client.HTTPConnection("127.0.0.1", httpd.server_address[1]) + + def _ready_state(self, runner: FakeRunner | None = None) -> ServerState: + state = ServerState() + # The loader thread resolves fps before flipping to "ready"; mirror that here. + state.fps = 20.0 + state.runner = runner or FakeRunner() + state.status = "ready" + return state + + def _auth_header(self) -> dict: + return {"Authorization": f"Bearer {self.TEST_KEY}"} + + def _infer(self, conn: http.client.HTTPConnection, payload: dict): + conn.request( + "POST", + "/infer", + body=json.dumps(payload).encode(), + headers=self._auth_header(), + ) + resp = conn.getresponse() + return resp.status, json.loads(resp.read()) + + def _valid_payload(self) -> dict: + blank = encode_bgr_jpeg_b64(np.zeros((4, 4, 3), dtype=np.uint8)) + # Request image names are the checkpoint's own camera keys; FakeRunner + # expects "scene", so the observation supplies "scene". + return {"state": [0.0, 0.0], "task": "stack", "images": {"scene": blank}} + + def test_handler_sets_connection_timeout(self) -> None: + """A half-open connection cannot park its handler thread forever: the + handler applies a socket timeout to every connection.""" + handler_cls = make_handler(ServerState()) + self.assertEqual(handler_cls.timeout, REQUEST_SOCKET_TIMEOUT_SECONDS) + self.assertGreater(REQUEST_SOCKET_TIMEOUT_SECONDS, 0) + + def test_health_reports_loading(self) -> None: + """GET /health during model load reports 'loading', usable as a startup probe.""" + conn = self._start(ServerState()) + conn.request("GET", "/health") + resp = conn.getresponse() + + self.assertEqual(resp.status, 200) + self.assertEqual(json.loads(resp.read())["status"], "loading") + + def test_health_reports_error_with_detail(self) -> None: + """GET /health after a failed load carries the load error for diagnosis.""" + state = ServerState() + state.status = "error" + state.detail = "ValueError: no config.json" + conn = self._start(state) + conn.request("GET", "/health") + body = json.loads(conn.getresponse().read()) + + self.assertEqual(body["status"], "error") + self.assertIn("no config.json", body["detail"]) + + def test_unknown_path_is_404(self) -> None: + """A request to any path other than /health or /infer is rejected, not routed.""" + conn = self._start(self._ready_state()) + conn.request("GET", "/unknown") + self.assertEqual(conn.getresponse().status, 404) + + def test_infer_while_loading_is_503_with_message(self) -> None: + """POST /infer during model load answers 503 'still loading', which the + adapter relays verbatim to the MoveIt Pro UI.""" + conn = self._start(ServerState()) + status, body = self._infer(conn, {"task": "x"}) + + self.assertEqual(status, 503) + self.assertIn("still loading", body["error"]) + + def test_infer_after_failed_load_is_500_with_detail(self) -> None: + """POST /infer after a failed load relays the load error, not a generic 500.""" + state = ServerState() + state.status = "error" + state.detail = "ValueError: checkpoint directory does not exist" + conn = self._start(state) + status, body = self._infer(conn, {"task": "x"}) + + self.assertEqual(status, 500) + self.assertIn("checkpoint directory does not exist", body["error"]) + + def test_infer_without_token_is_401(self) -> None: + """POST /infer without the shared key is rejected before any other check.""" + conn = self._start(self._ready_state()) + conn.request("POST", "/infer", body=b"{}") + resp = conn.getresponse() + + self.assertEqual(resp.status, 401) + self.assertIn("MOVEIT_FRONTEND_KEY", json.loads(resp.read())["error"]) + + def test_infer_with_wrong_token_is_401(self) -> None: + """A mismatched key is rejected the same as a missing one.""" + conn = self._start(self._ready_state()) + conn.request( + "POST", + "/infer", + body=b"{}", + headers={"Authorization": "Bearer wrong-key"}, + ) + self.assertEqual(conn.getresponse().status, 401) + + def test_infer_while_loading_still_requires_token(self) -> None: + """Auth wraps the whole endpoint: an unauthenticated probe cannot even + distinguish the loading state.""" + conn = self._start(ServerState()) + conn.request("POST", "/infer", body=b"{}") + self.assertEqual(conn.getresponse().status, 401) + + def test_infer_with_non_ascii_token_is_401(self) -> None: + """A non-ASCII token gets a clean 401, not a dropped connection: + compare_digest on str raises TypeError for non-ASCII input.""" + conn = self._start(self._ready_state()) + conn.request( + "POST", + "/infer", + body=b"{}", + headers={"Authorization": "Bearer café-key"}, + ) + self.assertEqual(conn.getresponse().status, 401) + + def test_infer_accepts_case_insensitive_bearer_scheme(self) -> None: + """The auth scheme is case-insensitive per RFC 7235.""" + conn = self._start(self._ready_state()) + conn.request( + "POST", + "/infer", + body=json.dumps(self._valid_payload()).encode(), + headers={"Authorization": f"bEaReR {self.TEST_KEY}"}, + ) + self.assertEqual(conn.getresponse().status, 200) + + def test_health_needs_no_token(self) -> None: + """GET /health stays token-free so health probes keep working.""" + conn = self._start(self._ready_state()) + conn.request("GET", "/health") + resp = conn.getresponse() + + self.assertEqual(resp.status, 200) + self.assertEqual(json.loads(resp.read())["status"], "ready") + + def test_infer_malformed_json_is_400(self) -> None: + """A body that isn't valid JSON is rejected before it reaches the policy.""" + conn = self._start(self._ready_state()) + conn.request("POST", "/infer", body=b"not json", headers=self._auth_header()) + self.assertEqual(conn.getresponse().status, 400) + + def test_infer_negative_content_length_is_400(self) -> None: + """A negative Content-Length is rejected before any body read.""" + conn = self._start(self._ready_state()) + conn.putrequest("POST", "/infer", skip_accept_encoding=True) + conn.putheader("Authorization", f"Bearer {self.TEST_KEY}") + conn.putheader("Content-Length", "-1") + conn.endheaders() + self.assertEqual(conn.getresponse().status, 400) + + def test_infer_oversized_body_is_413(self) -> None: + """A declared body size over the ceiling is rejected before it is read.""" + conn = self._start(self._ready_state()) + conn.request( + "POST", + "/infer", + body=b"x", + headers={"Content-Length": str(2**40), **self._auth_header()}, + ) + self.assertEqual(conn.getresponse().status, 413) + + def test_infer_success_returns_chunk_and_dt(self) -> None: + """A valid POST /infer runs the policy and returns the chunk with dt=1/fps.""" + conn = self._start(self._ready_state()) + status, body = self._infer(conn, self._valid_payload()) + + self.assertEqual(status, 200) + self.assertEqual(body["action_chunk"], [[0.1, 0.2]]) + self.assertEqual(body["action_chunk_raw"], [[0.5, 0.5]]) + self.assertAlmostEqual(body["dt"], 0.05) + + def test_infer_missing_expected_camera_is_400_naming_it(self) -> None: + """Images that omit one of the checkpoint's cameras are rejected as the + caller's error (400), because lerobot would otherwise zero-fill the + camera and run the policy blind.""" + conn = self._start(self._ready_state()) # FakeRunner expects "scene" + blank = encode_bgr_jpeg_b64(np.zeros((4, 4, 3), dtype=np.uint8)) + # The request supplies "front", not the checkpoint's expected "scene". + status, body = self._infer( + conn, {"state": [0.0, 0.0], "task": "stack", "images": {"front": blank}} + ) + + self.assertEqual(status, 400) + self.assertIn("scene", body["error"]) + self.assertIn("image_names", body["error"]) + + def test_infer_native_camera_names_accepted(self) -> None: + """A complete set of dataset-native names serves a chunk: the checkpoint's + own rename step maps them onto the model slots.""" + runner = FakeRunner(native_map={"front": "scene"}) + conn = self._start(self._ready_state(runner)) + blank = encode_bgr_jpeg_b64(np.zeros((4, 4, 3), dtype=np.uint8)) + status, body = self._infer( + conn, {"state": [0.0, 0.0], "task": "stack", "images": {"front": blank}} + ) + + self.assertEqual(status, 200) + self.assertIn("action_chunk", body) + + def test_infer_slot_names_on_renaming_checkpoint_rejected_with_fix(self) -> None: + """A renaming checkpoint takes its dataset camera names only; the + config.json slot names are refused with a message naming the names to + use instead.""" + runner = FakeRunner(native_map={"front": "scene"}) + conn = self._start(self._ready_state(runner)) + blank = encode_bgr_jpeg_b64(np.zeros((4, 4, 3), dtype=np.uint8)) + status, body = self._infer( + conn, {"state": [0.0, 0.0], "task": "stack", "images": {"scene": blank}} + ) + + self.assertEqual(status, 400) + self.assertIn("front", body["error"]) + self.assertIn("image_names", body["error"]) + + def test_infer_partial_rename_native_full_set_accepted(self) -> None: + """A checkpoint renaming only some cameras accepts the mixed native set: + the dataset name where one exists, the slot name where it does not.""" + runner = FakeRunner(camera_keys=["scene", "aux"], native_map={"front": "scene"}) + conn = self._start(self._ready_state(runner)) + blank = encode_bgr_jpeg_b64(np.zeros((4, 4, 3), dtype=np.uint8)) + status, body = self._infer( + conn, + { + "state": [0.0, 0.0], + "task": "stack", + "images": {"front": blank, "aux": blank}, + }, + ) + + self.assertEqual(status, 200) + self.assertIn("action_chunk", body) + + def test_infer_extra_camera_name_is_400_naming_it(self) -> None: + """A request carrying both a dataset name and the slot it renames to is + refused: lerobot's rename step would silently overwrite one with the + other and feed the policy the wrong camera.""" + runner = FakeRunner(native_map={"front": "scene"}) + conn = self._start(self._ready_state(runner)) + blank = encode_bgr_jpeg_b64(np.zeros((4, 4, 3), dtype=np.uint8)) + status, body = self._infer( + conn, + { + "state": [0.0, 0.0], + "task": "stack", + "images": {"front": blank, "scene": blank}, + }, + ) + + self.assertEqual(status, 400) + self.assertIn("unexpected", body["error"]) + self.assertIn("scene", body["error"]) + + def test_infer_partial_rename_missing_unrenamed_camera_is_400(self) -> None: + """Covering only the renamed camera is refused: the unrenamed one would + be silently zero-filled and the policy would run partially blind.""" + runner = FakeRunner(camera_keys=["scene", "aux"], native_map={"front": "scene"}) + conn = self._start(self._ready_state(runner)) + blank = encode_bgr_jpeg_b64(np.zeros((4, 4, 3), dtype=np.uint8)) + status, body = self._infer( + conn, {"state": [0.0, 0.0], "task": "stack", "images": {"front": blank}} + ) + + self.assertEqual(status, 400) + self.assertIn("aux", body["error"]) + + def test_infer_state_width_mismatch_is_400_with_both_widths(self) -> None: + """A state narrower than the checkpoint expects is the caller's error + (400) and names both widths.""" + conn = self._start(self._ready_state()) + payload = self._valid_payload() + payload["state"] = [0.0] + status, body = self._infer(conn, payload) + + self.assertEqual(status, 400) + self.assertIn("1-dim state", body["error"]) + self.assertIn("expects 2", body["error"]) + + def test_infer_non_list_state_is_400(self) -> None: + """A scalar or string state is the caller's error (400), not an + internal 500 from torch.tensor.""" + conn = self._start(self._ready_state()) + payload = self._valid_payload() + payload["state"] = "oops" + status, body = self._infer(conn, payload) + + self.assertEqual(status, 400) + self.assertIn("must be a list", body["error"]) + + def test_infer_non_string_image_value_is_400(self) -> None: + """A non-string camera value is the caller's error (400), not a + TypeError-turned-500 from base64.""" + conn = self._start(self._ready_state()) + payload = self._valid_payload() + payload["images"]["scene"] = 5 + status, body = self._infer(conn, payload) + + self.assertEqual(status, 400) + self.assertIn("base64", body["error"]) + + def test_infer_non_finite_state_is_400(self) -> None: + """A NaN joint position from a degraded publisher is the caller's + error (400), not policy input.""" + conn = self._start(self._ready_state()) + payload = self._valid_payload() + payload["state"] = [0.0, float("nan")] + status, body = self._infer(conn, payload) + + self.assertEqual(status, 400) + self.assertIn("finite", body["error"]) + + def test_infer_non_finite_prev_chunk_is_400(self) -> None: + """A NaN in the RTC carryover is the caller's error (400), not + guidance input.""" + conn = self._start(self._ready_state()) + payload = self._valid_payload() + payload["prev_chunk_left_over"] = [[0.1, float("nan")]] + status, body = self._infer(conn, payload) + + self.assertEqual(status, 400) + self.assertIn("prev_chunk_left_over", body["error"]) + + def test_infer_exception_is_500_with_error_field(self) -> None: + """A policy exception returns 500 with {"error": ...}; the adapter parses the + body before checking the status, so the detail still reaches the operator.""" + conn = self._start(self._ready_state(FakeRunner(RuntimeError("cuda OOM")))) + status, body = self._infer(conn, self._valid_payload()) + + self.assertEqual(status, 500) + self.assertIn("cuda OOM", body["error"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/vla_sim/docker/vla_inference_server.py b/src/vla_sim/docker/vla_inference_server.py new file mode 100644 index 000000000..b3ca3c292 --- /dev/null +++ b/src/vla_sim/docker/vla_inference_server.py @@ -0,0 +1,931 @@ +#!/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. + +"""LeRobot inference server for MoveIt Pro's ExecutePolicy. + +Serves POST /infer and GET /health over HTTP. Runs in its own container (see +Dockerfile.vla_inference_server and the workspace docker-compose.yaml +`inference_server` service) so torch/lerobot stay out of the MoveIt Pro +images; the in-config adapter node (script/get_action_chunk_adapter.py) +bridges the /get_action_chunk ROS service to this server. + +/infer requires the deployment's shared MOVEIT_FRONTEND_KEY as an +`Authorization: Bearer` token, the same key the MoveIt Pro web backend +endpoints use; a blank or unset key parks the server in the error state (fail +closed, matching those endpoints). /health stays token-free for health +probes. For a bare development run, export the documented dev key first +(`MOVEIT_FRONTEND_KEY=moveit-secret-key`). + +The socket binds before the checkpoint loads: /health reports +loading|ready|error and /infer answers 503 (loading) or 500 (load failed) with +the same detail until the model is ready, so the adapter can tell the operator +exactly what is wrong from the MoveIt Pro UI. A bad or missing checkpoint +parks the server in the error state, surfacing the problem through the +service instead of exiting into a compose restart loop. + +Each knob resolves as: an explicit CLI flag > the per-config model-serving +YAML (vla_serving.yaml, default /vla_config/vla_serving.yaml, mounted from +src/vla_sim/config/) > a built-in default. A missing or empty YAML file is +fine, so a bare `python vla_inference_server.py --checkpoint ` +still works for development; a malformed file or value parks the server in +the error state. +""" + +import argparse +import base64 +import hmac +import json +import math +import os +import threading +import time +import traceback +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path + +import cv2 +import numpy as np +import torch +import yaml +from huggingface_hub import hf_hub_download +from huggingface_hub.errors import GatedRepoError, RepositoryNotFoundError + +from lerobot.configs.types import RTCAttentionSchedule +from lerobot.policies.factory import get_policy_class, make_pre_post_processors +from lerobot.policies.rtc.configuration_rtc import RTCConfig + +# pi0.5 checkpoints save a processor pipeline that references +# 'relative_actions_processor', an alias lerobot does not always auto-register; +# without it make_pre_post_processors raises +# "Processor step 'relative_actions_processor' not found". +from lerobot.processor import ProcessorStepRegistry +from lerobot.processor.relative_action_processor import RelativeActionsProcessorStep + +try: + ProcessorStepRegistry.get("relative_actions_processor") +except Exception: + ProcessorStepRegistry.register("relative_actions_processor")( + RelativeActionsProcessorStep + ) + +TESTED_POLICY_TYPES = ("smolvla", "pi05") + +# Generous ceiling over a multi-camera base64 observation; bounds the memory +# one connection can demand before the body is read. +MAX_BODY_BYTES = 32 * 1024 * 1024 + +# Per-connection socket timeout. Each connection gets a handler thread, so +# without a timeout a client that opens a connection and never completes its +# request parks that thread forever. Generous over the largest loopback body +# read; inference time is not affected (no socket reads happen during it). +REQUEST_SOCKET_TIMEOUT_SECONDS = 30 + +# The per-config model-serving YAML, mounted read-only from the workspace's +# src/vla_sim/config/. Overridable with --config for a standalone `docker run`. +DEFAULT_CONFIG_PATH = "/vla_config/vla_serving.yaml" + + +def load_serving_config(path: str) -> dict: + """Read the per-config model-serving YAML into a dict of knob values. + + A missing or empty file returns {}: built-in defaults still apply, so a + bare `docker run` needs no config. A malformed file (or one whose top + level is not a mapping) raises, so a typo in the operator's tuning surface + fails loudly instead of silently serving with the wrong knobs. + """ + file = Path(path).expanduser() + if not file.is_file(): + log(f"no serving config at '{path}'; using built-in defaults") + return {} + loaded = yaml.safe_load(file.read_text()) + if loaded is None: + log(f"serving config '{path}' is empty; using built-in defaults") + return {} + if not isinstance(loaded, dict): + raise ValueError( + f"serving config '{path}' must be a YAML mapping of knob names to " + f"values, not a {type(loaded).__name__}" + ) + return loaded + + +def resolve_default(yaml_value, builtin): + """Pick an argparse default: the YAML value wins when present. + + argparse layers an explicit CLI flag on top of this, giving the full + precedence CLI flag > YAML > built-in. A YAML value of 0 or "" is honored, + since 0 is a meaningful "auto" sentinel for fps and state_dim. + """ + return builtin if yaml_value is None else yaml_value + + +def load_checkpoint_file(checkpoint: str, filename: str) -> dict: + """Read a JSON file from a local checkpoint directory or an HF repo. + + A local path wins when it exists; anything else must look like an HF repo + id ("org/name"), fetched through the cache (HF_HUB_OFFLINE and HF_TOKEN + apply). + """ + path = Path(checkpoint).expanduser() + if path.is_dir(): + file = path / filename + if not file.is_file(): + raise FileNotFoundError(f"'{path}' has no {filename}") + return json.loads(file.read_text()) + if "/" not in checkpoint: + raise ValueError( + f"checkpoint '{checkpoint}' is neither a local directory nor a " + "Hugging Face repo id" + ) + return json.loads(Path(hf_hub_download(checkpoint, filename)).read_text()) + + +def resolve_policy_type(checkpoint: str, override: str) -> str: + """Read the policy family from the checkpoint's config.json unless overridden.""" + if override: + return override + policy_type = load_checkpoint_file(checkpoint, "config.json").get("type", "") + if not policy_type: + raise ValueError( + f"the config.json of '{checkpoint}' carries no 'type' field; " + "set policy_class in vla_serving.yaml (or --policy-class) " + "explicitly" + ) + return policy_type + + +def resolve_fps(checkpoint: str, fps: float) -> float: + """Resolve the policy's training rate, preferring the explicit value. + + The chunk is played at 1/fps seconds per step; a wrong value scales every + commanded joint velocity, so an unresolvable rate is a startup error, + never a silent default. + """ + if fps > 0.0: + return fps + try: + config = load_checkpoint_file(checkpoint, "train_config.json") + except (GatedRepoError, RepositoryNotFoundError): + # The tailored HF-access advice in load_policy beats a generic + # missing-fps message. + raise + except Exception: + config = {} + from_config = config.get("dataset", {}).get("fps") or config.get("fps") + if from_config and float(from_config) > 0: + return float(from_config) + raise ValueError( + f"could not read the training fps from the train_config.json of " + f"'{checkpoint}'; set fps in vla_serving.yaml (or --fps) to the " + "rate the policy was trained at" + ) + + +def resolve_device(requested: str, cuda_available: bool) -> str: + """Resolve the torch device, failing loudly when an explicit request can't be honored. + + 'auto' picks cuda when torch reports a usable GPU and falls back to cpu. + An explicit cuda request on a host without one is a startup error, never a + silent cpu fallback, so pacing tuned for a GPU cannot quietly run an order + of magnitude slower. + """ + if requested == "auto": + return "cuda" if cuda_available else "cpu" + if requested.startswith("cuda") and not cuda_available: + raise ValueError( + f"device '{requested}' was requested but this torch build reports " + "no usable GPU; run the container with the NVIDIA runtime " + "(GPU serving is automatic on NVIDIA machines under the launcher) " + "or set device: auto in vla_serving.yaml" + ) + return requested + + +def resolve_rtc_horizon( + inference_delay: int, guidance_horizon: int, default_guidance: int +) -> int: + """Map the service's soft-guidance width onto lerobot's RTC horizon. + + lerobot's execution_horizon is the end index of the guided region measured + from the chunk start (get_prefix_weights(start=inference_delay, + end=execution_horizon)), while the GetActionChunk contract's + guidance_horizon is that region's width past the frozen prefix, zero + deferring to the server default. Passing the width through unconverted + would shrink the frozen prefix whenever the width is smaller than the + inference delay. + """ + width = guidance_horizon if guidance_horizon > 0 else default_guidance + return inference_delay + width + + +def resolve_rtc_schedule(name: str) -> RTCAttentionSchedule: + """Map the rtc_schedule knob onto lerobot's enum, naming the valid values on a typo.""" + try: + return RTCAttentionSchedule[name] + except KeyError: + valid = ", ".join(schedule.name for schedule in RTCAttentionSchedule) + raise ValueError( + f"rtc_schedule '{name}' is not a known RTC schedule; set " + f"rtc_schedule in vla_serving.yaml (or pass --rtc-schedule) " + f"to one of: {valid}" + ) from None + + +def decode_image_b64(data: str) -> torch.Tensor: + """base64 JPEG -> CHW float32 [0,1] RGB tensor.""" + buf = np.frombuffer(base64.b64decode(data), dtype=np.uint8) + bgr = cv2.imdecode(buf, cv2.IMREAD_COLOR) + if bgr is None: + raise ValueError("cv2.imdecode failed on /infer image") + rgb = np.ascontiguousarray(bgr[:, :, ::-1]) + return torch.from_numpy(rgb).float().permute(2, 0, 1) / 255.0 + + +def native_camera_map(steps: list) -> dict: + """Map the checkpoint's dataset-native camera names to its model slot names. + + Read from the first preprocessor step whose rename mapping touches image + keys. The mapping is recognized by its shape (a rename_map dict) rather + than by a lerobot class, so a lerobot relayout degrades to the model slot + names instead of crashing the server before the socket binds. Empty when + no step renames images (the dataset already used the model's slot names). + """ + prefix = "observation.images." + image_maps = [] + for step in steps: + rename_map = getattr(step, "rename_map", None) + if not isinstance(rename_map, dict): + continue + image_renames = { + src.removeprefix(prefix): dst.removeprefix(prefix) + for src, dst in rename_map.items() + if src.startswith(prefix) and dst.startswith(prefix) + } + if image_renames: + image_maps.append(image_renames) + if len(image_maps) > 1: + log( + f"WARNING: {len(image_maps)} preprocessor steps rename cameras; " + "the request names are derived from the first" + ) + return image_maps[0] if image_maps else {} + + +def request_camera_names(slot_keys: list, native_map: dict) -> list: + """The camera names an /infer request must carry, in checkpoint order. + + The dataset-native name where the checkpoint's rename step defines one, the + model slot name for a camera it does not rename. + """ + slot_to_native = {slot: native for native, slot in native_map.items()} + if len(slot_to_native) < len(native_map): + log( + "WARNING: the checkpoint's rename map sends several dataset camera " + "names to the same model slot; requests must use the last one" + ) + return [slot_to_native.get(slot, slot) for slot in slot_keys] + + +class PolicyRunner: + """Owns the loaded policy and serializes inference calls. + + Loading passes policy_cfg by keyword and overrides the device on both + processors, which merged pi0.5 checkpoints need: their config declares a + padded 32-dim state while the saved normalizer stats carry the trained + width. + """ + + def __init__( + self, + checkpoint: str, + policy_type: str, + device: str, + guidance_horizon: int, + rtc_schedule: str, + state_dim: int, + ): + self.device = device + self.state_dim = state_dim + self.guidance_horizon = guidance_horizon + self.lock = threading.Lock() + + # Resolve before the slow checkpoint load so a schedule typo fails fast. + schedule = resolve_rtc_schedule(rtc_schedule) + + self.policy = get_policy_class(policy_type).from_pretrained(checkpoint) + self.policy.to(device) + self.policy.eval() + + # infer() passes the horizon per call on every RTC request, so the + # config's own execution_horizon never applies; only the enable and + # schedule matter here. + self.policy.config.rtc_config = RTCConfig( + enabled=True, + prefix_attention_schedule=schedule, + ) + self.policy.init_rtc_processor() + + self.pre, self.post = make_pre_post_processors( + policy_cfg=self.policy.config, + pretrained_path=checkpoint, + preprocessor_overrides={"device_processor": {"device": device}}, + postprocessor_overrides={"device_processor": {"device": device}}, + ) + + # Derived once here, on the loader thread before the runner is + # published, so the load log and every /infer validation report the + # same set. + self.request_names = request_camera_names( + self.expected_camera_keys(), native_camera_map(self.pre.steps) + ) + + def expected_state_dim(self) -> int: + # The state width config.json declares is not reliable: a checkpoint + # can declare its pretraining base's width or its architecture's padded + # maximum while the saved normalizer stats carry the width it was + # actually trained on. state_dim is the caller asserting that trained + # width; unset, the declared value is used and a mismatched checkpoint + # fails at warmup. + if self.state_dim > 0: + return self.state_dim + return int(self.policy.config.input_features["observation.state"].shape[0]) + + def expected_camera_keys(self) -> list: + """The camera keys the checkpoint was trained on, without the feature prefix.""" + return [ + key.removeprefix("observation.images.") + for key in self.policy.config.input_features + if "image" in key + ] + + @torch.no_grad() + def infer( + self, + images: dict, + state: list, + prompt: str, + prev_chunk: np.ndarray | None, + inference_delay: int, + guidance_horizon: int, + ) -> tuple[np.ndarray, np.ndarray]: + """Run one inference. Returns (absolute_actions, normalized_actions), both (T, A).""" + obs = { + f"observation.images.{key}": tensor.to(self.device) + for key, tensor in images.items() + } + obs["observation.state"] = torch.tensor( + state, dtype=torch.float32, device=self.device + ) + obs["task"] = prompt + + kwargs = {} + if prev_chunk is not None and prev_chunk.size > 0: + kwargs["prev_chunk_left_over"] = torch.tensor( + prev_chunk, dtype=torch.float32, device=self.device + ) + kwargs["inference_delay"] = inference_delay + kwargs["execution_horizon"] = resolve_rtc_horizon( + inference_delay, guidance_horizon, self.guidance_horizon + ) + + with self.lock: + self.policy.reset() + chunk = self.policy.predict_action_chunk( + self.pre(obs), **kwargs + ) # (1, T, A) normalized + normalized = chunk.squeeze(0).detach().cpu().numpy() + steps = [self.post(chunk[:, i, :]) for i in range(chunk.shape[1])] + absolute = torch.stack(steps, dim=1).squeeze(0).detach().cpu().numpy() + return absolute, normalized + + def warmup(self) -> tuple[float, float, int]: + """Full-size warmup so the first real chunk does not pay model compile/cache costs. + + Returns (cold_s, steady_s, chunk_steps): the first pass carries the + one-time compile/cache costs, the second approximates the per-request + latency that execution pacing must absorb. + """ + images = {key: torch.zeros(3, 224, 224) for key in self.expected_camera_keys()} + state = [0.0] * self.expected_state_dim() + start = time.perf_counter() + chunk, _ = self.infer(images, state, "warmup", None, 0, 0) + cold_s = time.perf_counter() - start + start = time.perf_counter() + self.infer(images, state, "warmup", None, 0, 0) + steady_s = time.perf_counter() - start + return cold_s, steady_s, chunk.shape[0] + + +class ServerState: + """Load status shared between the loader thread and the HTTP handlers.""" + + def __init__(self): + self.status = "loading" + self.detail = "" + self.runner: PolicyRunner | None = None + # Resolved by the loader thread; meaningful once status is "ready". + self.fps = 0.0 + # Shared secret /infer requests must present; set from + # MOVEIT_FRONTEND_KEY in main() before serve_forever() accepts any + # request. + self.frontend_key = "" + + +def log(message: str) -> None: + print(f"[vla_inference_server] {message}", flush=True) + + +def hub_access_error_message(checkpoint: str, gated: bool, token_present: bool) -> str: + """Actionable message for a Hugging Face download rejected for access reasons. + + The rejection can come from the checkpoint itself or from a gated + dependency it resolves (the pi0.5 processor pulls the gated + google/paligemma tokenizer). Whether HF_TOKEN is present decides the + advice; the token value is never echoed. + """ + if gated and token_present: + return ( + f"a Hugging Face repo needed by checkpoint '{checkpoint}' is gated " + "and the HF_TOKEN account has not been granted access: accept the " + "model's license on huggingface.co with that account, then restart" + ) + if gated: + return ( + f"a Hugging Face repo needed by checkpoint '{checkpoint}' is gated " + "and HF_TOKEN is not set in the environment: export HF_TOKEN with " + "a token from an account that accepted the model's license, then " + "restart" + ) + if token_present: + return ( + f"checkpoint '{checkpoint}' was not found on Hugging Face with the " + "provided HF_TOKEN: check the checkpoint name in vla_serving.yaml " + "and that the token's account can access the repo" + ) + return ( + f"checkpoint '{checkpoint}' was not found on Hugging Face: check the " + "checkpoint name in vla_serving.yaml; a private repo also needs " + "HF_TOKEN exported in the environment" + ) + + +def load_policy(state: ServerState, args: argparse.Namespace) -> None: + """Load + warm the policy in the background; on failure park in the error state.""" + try: + # A malformed serving config was deferred out of parse_args so the + # socket could bind first; surface it here like any other load failure. + if args.config_error: + raise ValueError(f"could not read the serving config: {args.config_error}") + # A missing checkpoint parks in the error state like any other config + # problem: no restart loop can supply the argument, and /health plus + # the objective's UI messages then name the fix. + if not args.checkpoint: + raise ValueError( + "set the checkpoint in vla_serving.yaml (or --checkpoint) to " + "a local LeRobot checkpoint directory or an HF repo id" + ) + # Resolving the fps can read the checkpoint's train_config.json (a + # download for hub checkpoints), so it happens here rather than before + # the socket binds, and an unresolvable rate parks in the error state + # instead of exiting into a compose restart loop. + state.fps = resolve_fps(args.checkpoint, args.fps) + policy_type = resolve_policy_type(args.checkpoint, args.policy_class) + device = resolve_device(args.device, torch.cuda.is_available()) + if policy_type not in TESTED_POLICY_TYPES: + log( + f"WARNING: policy family '{policy_type}' is untested with this " + f"server (tested: {', '.join(TESTED_POLICY_TYPES)}); loading best-effort" + ) + log( + f"loading {policy_type} checkpoint '{args.checkpoint}' on '{device}' " + f"(torch {torch.__version__}) ..." + ) + runner = PolicyRunner( + args.checkpoint, + policy_type, + device, + args.guidance_horizon, + args.rtc_schedule, + args.state_dim, + ) + + image_features = [ + k for k in runner.policy.config.input_features if "image" in k + ] + log( + f"checkpoint expects {len(image_features)} camera(s) {image_features} " + f"(request names: {runner.request_names}) and a " + f"{runner.expected_state_dim()}-dim state; chunk plays at " + f"dt={1.0 / state.fps:.4f}s" + ) + + # Warmup pre-pays model compile/cache costs; a failure here is logged, + # not fatal, because a padded-state pi0.5 config can reject the blank + # observation while real requests, which carry the true shape, still + # work (set state_dim in vla_serving.yaml to warm up cleanly). + try: + cold_s, steady_s, chunk_steps = runner.warmup() + # Real-time chunking stays feasible only while one inference fits + # into half a chunk of playback time: each seam must commit at + # least latency/dt steps yet leave at least as many uncommitted, + # capping tolerable latency at (chunk/2)*dt. + budget_s = (chunk_steps / 2.0) / state.fps + if steady_s > budget_s: + log( + f"WARNING: inference takes {steady_s:.2f}s per " + f"{chunk_steps}-step chunk on '{device}', over the " + f"{budget_s:.2f}s real-time budget at {state.fps:g} fps; " + "execution will starve at chunk seams. Serve on a faster " + "device (the launcher uses the GPU automatically on NVIDIA " + "machines) or use a policy this machine can serve in time. The " + "objective's committed_action_steps x dt sets the " + "tighter per-run budget." + ) + else: + log( + f"warmup done: {steady_s:.2f}s per {chunk_steps}-step chunk " + f"(cold start {cold_s:.2f}s), within the {budget_s:.2f}s " + f"real-time budget at {state.fps:g} fps" + ) + except Exception as exc: + log( + f"WARNING: warmup inference failed ({type(exc).__name__}: {exc}); " + "continuing, the first request pays the cold cost" + ) + + state.runner = runner + state.status = "ready" + log("ready") + except (GatedRepoError, RepositoryNotFoundError) as exc: + traceback.print_exc() + state.detail = hub_access_error_message( + args.checkpoint, + isinstance(exc, GatedRepoError), + bool(os.environ.get("HF_TOKEN")), + ) + state.status = "error" + log(f"FATAL: model load failed: {state.detail}") + except Exception as exc: + traceback.print_exc() + state.detail = f"{type(exc).__name__}: {exc}" + state.status = "error" + log(f"FATAL: model load failed: {state.detail}") + + +def run_inference(state: ServerState, payload: dict) -> dict: + """Validate one /infer payload and run it through the loaded policy. + + new_episode in the payload is informational only: PolicyRunner resets per + call, and the episode boundary is carried by an empty prev_chunk. + """ + for key in ("state", "images", "task"): + if key not in payload: + raise ValueError(f"/infer payload is missing '{key}'") + if not isinstance(payload["images"], dict) or not all( + isinstance(v, str) for v in payload["images"].values() + ): + raise ValueError( + "/infer payload 'images' must map camera names to base64-encoded JPEG strings" + ) + if not isinstance(payload["state"], list) or not all( + isinstance(v, (int, float)) and math.isfinite(v) for v in payload["state"] + ): + raise ValueError( + "/infer payload 'state' must be a list of finite joint positions" + ) + runner = state.runner + expected_state = runner.expected_state_dim() + robot_state = payload["state"] + if len(robot_state) != expected_state: + raise ValueError( + f"request carries a {len(robot_state)}-dim state but the " + f"checkpoint expects {expected_state} dims" + ) + + # Request image names are the checkpoint's own camera names: the dataset + # names baked into its preprocessor rename step, or, for a camera the + # checkpoint does not rename, its config.json slot name. Exactly that set + # is required: lerobot zero-fills an expected camera the observation + # lacks, and its rename step lets an unexpected name silently overwrite a + # renamed camera's slot; either way the policy would run on wrong images, + # so refuse before decoding anything. + expected_names = runner.request_names + missing = [name for name in expected_names if name not in payload["images"]] + unexpected = sorted(payload["images"].keys() - set(expected_names)) + if missing or unexpected: + problems = [] + if missing: + problems.append(f"is missing camera(s) {missing}") + if unexpected: + problems.append(f"carries unexpected camera(s) {unexpected}") + raise ValueError( + f"the request {' and '.join(problems)} but this checkpoint takes " + f"exactly {expected_names}; set each of the objective's " + "image_names to the checkpoint's name for the camera on the " + "matching image_topics entry" + ) + images = {name: decode_image_b64(data) for name, data in payload["images"].items()} + + prev_chunk = None + prev = payload.get("prev_chunk_left_over") + if prev: + try: + prev_chunk = np.asarray(prev, dtype=float) + except (TypeError, ValueError) as exc: + raise ValueError( + f"/infer payload 'prev_chunk_left_over' is not a numeric " + f"array ({exc})" + ) from exc + if prev_chunk.ndim != 2 or not np.isfinite(prev_chunk).all(): + raise ValueError( + "/infer payload 'prev_chunk_left_over' must be a 2-D array of " + "finite numbers" + ) + inference_delay = int(payload.get("inference_delay", 0)) + guidance_horizon = int(payload.get("guidance_horizon", 0)) + if inference_delay < 0 or guidance_horizon < 0: + raise ValueError( + "/infer payload 'inference_delay' and 'guidance_horizon' must be " + "non-negative" + ) + + absolute, normalized = runner.infer( + images, + robot_state, + payload["task"], + prev_chunk, + inference_delay, + guidance_horizon, + ) + return { + "action_chunk": absolute.tolist(), + "action_chunk_raw": normalized.tolist(), + "dt": 1.0 / state.fps, + } + + +def make_handler(state: ServerState): + class Handler(BaseHTTPRequestHandler): + # Applied to the connection socket by the base class, so a stalled + # request read raises and frees the thread instead of parking it. + timeout = REQUEST_SOCKET_TIMEOUT_SECONDS + + def _send(self, code: int, obj: dict) -> None: + body = json.dumps(obj).encode() + try: + self.send_response(code) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + except (BrokenPipeError, ConnectionResetError): + # The client gave up waiting (its timeout is shorter than this + # inference took); one line beats a stack trace per abort. + log(f"client disconnected before the {code} response was sent") + + def do_GET(self): + if self.path != "/health": + self._send(404, {"error": "not found"}) + return + health = {"status": state.status} + if state.status == "error": + health["detail"] = state.detail + elif state.status == "ready": + health["device"] = state.runner.device + self._send(200, health) + + def _authorized(self) -> bool: + # Same contract as the MoveIt Pro REST auth middleware: the shared + # key as an `Authorization: Bearer` token, compared in constant + # time. /health never reaches this check. + header = self.headers.get("Authorization", "") + scheme, _, token = header.partition(" ") + if scheme.lower() != "bearer" or not token.strip(): + return False + # Compare bytes: compare_digest raises TypeError on non-ASCII str, + # and header values arrive latin-1-decoded, so a crafted header + # would otherwise drop the connection instead of getting a 401. + return hmac.compare_digest( + token.strip().encode(), state.frontend_key.encode() + ) + + def do_POST(self): + if self.path != "/infer": + self._send(404, {"error": "not found"}) + return + if not self._authorized(): + self._send( + 401, + { + "error": "/infer requires the deployment's " + "MOVEIT_FRONTEND_KEY as an 'Authorization: Bearer' " + "token" + }, + ) + return + if state.status == "loading": + self._send( + 503, + { + "error": "the inference server is still loading " + "the model; try again shortly" + }, + ) + return + if state.status == "error": + self._send(500, {"error": f"model load failed: {state.detail}"}) + return + try: + length = int(self.headers.get("Content-Length", 0)) + if length < 0: + # A negative length would make rfile.read() unbounded. + self._send(400, {"error": "invalid negative Content-Length"}) + return + if length > MAX_BODY_BYTES: + self._send( + 413, + { + "error": f"request body of {length} bytes " + f"exceeds the {MAX_BODY_BYTES}-byte " + "limit" + }, + ) + return + payload = json.loads(self.rfile.read(length)) + except ValueError as exc: + self._send(400, {"error": f"bad request: {exc}"}) + return + try: + self._send(200, run_inference(state, payload)) + except ValueError as exc: + # Request-shape problems (missing camera, state-width mismatch, + # undecodable image) are the caller's error, not a server fault. + self._send(400, {"error": f"{type(exc).__name__}: {exc}"}) + except Exception as exc: + traceback.print_exc() + self._send(500, {"error": f"{type(exc).__name__}: {exc}"}) + + def log_message(self, *args): + pass # quiet; call counting is done adapter-side + + return Handler + + +def parse_args() -> argparse.Namespace: + """CLI options resolving CLI flag > vla_serving.yaml > built-in default.""" + # Resolve the config path first so the YAML can seed the other defaults; a + # bootstrap parser reads only that flag. + bootstrap = argparse.ArgumentParser(add_help=False) + bootstrap.add_argument("--config", default=DEFAULT_CONFIG_PATH) + config_path = bootstrap.parse_known_args()[0].config + + config: dict = {} + config_errors: list = [] + try: + config = load_serving_config(config_path) + except Exception as exc: + # A malformed config must park in the error state after the socket + # binds, not crash before it. Defer: build defaults from built-ins and + # hand the error to the loader thread via the namespace. + config_errors.append(f"{type(exc).__name__}: {exc}") + + def numeric_default(name: str, cast, builtin): + # A non-numeric YAML value (the likeliest typo on the tuning surface) + # must also park post-bind, not crash into a compose restart loop. + value = resolve_default(config.get(name), builtin) + try: + return cast(value) + except (TypeError, ValueError): + config_errors.append(f"{name}: '{value}' is not a number") + return builtin + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--config", + default=config_path, + help="per-config model-serving YAML; missing or empty " + "is fine, malformed parks the error state", + ) + parser.add_argument( + "--checkpoint", + default=str(resolve_default(config.get("checkpoint"), "")), + help="local LeRobot checkpoint directory or HF repo id", + ) + parser.add_argument( + "--policy-class", + default=str(resolve_default(config.get("policy_class"), "")), + help="lerobot policy family (pi05 | smolvla | ...); " + "default: the checkpoint's config.json 'type'", + ) + parser.add_argument( + "--fps", + type=float, + default=numeric_default("fps", float, 0.0), + help="training fps; response dt=1/fps; 0 reads the " + "checkpoint's train_config.json", + ) + parser.add_argument( + "--device", + default=str(resolve_default(config.get("device"), "auto")), + help="torch device: auto | cpu | cuda", + ) + parser.add_argument("--port", type=int, default=8973) + parser.add_argument( + "--state-dim", + type=int, + default=numeric_default("state_dim", int, 0), + help="trained observation.state width when the " + "checkpoint's config.json declares a padded one " + "(0 = trust config.json)", + ) + parser.add_argument( + "--guidance-horizon", + type=int, + default=numeric_default("guidance_horizon", int, 8), + help="RTC soft-guidance width in steps past the frozen " + "prefix, used when a request's guidance_horizon " + "is zero", + ) + parser.add_argument( + "--rtc-schedule", + default=str(resolve_default(config.get("rtc_schedule"), "EXP")), + help="RTC guidance-weight schedule: " + + " | ".join(schedule.name for schedule in RTCAttentionSchedule), + ) + args = parser.parse_args() + args.config_error = "; ".join(config_errors) + return args + + +def apply_frontend_key(state: ServerState, raw_key: str | None) -> bool: + """Set the /infer auth key; park in the error state when it is blank. + + Fails closed like the other MOVEIT_FRONTEND_KEY consumers, but parks + instead of exiting so /health names the fix rather than a compose restart + loop hiding it. + + @param state: The server state to receive the key or the error. + @param raw_key: The MOVEIT_FRONTEND_KEY environment value, or None. + @return: True when the key is usable and the model load may proceed. + """ + key = (raw_key or "").strip() + if key: + state.frontend_key = key + return True + # /health serves this text without a token, so it points at the docs + # rather than naming any key value. + state.detail = ( + "MOVEIT_FRONTEND_KEY is required: /infer authenticates with the " + "deployment's shared key. Set it in the environment (see the MoveIt " + "Pro endpoint authentication guide), then restart" + ) + state.status = "error" + return False + + +def main() -> None: + # Compose forwards HF_TOKEN as an empty string when the host never set it; + # drop it so the Hub client sees a genuinely absent token. + if os.environ.get("HF_TOKEN") == "": + del os.environ["HF_TOKEN"] + args = parse_args() + state = ServerState() + httpd = ThreadingHTTPServer(("0.0.0.0", args.port), make_handler(state)) + if apply_frontend_key(state, os.environ.get("MOVEIT_FRONTEND_KEY")): + threading.Thread(target=load_policy, args=(state, args), daemon=True).start() + log(f"listening on 0.0.0.0:{args.port}; loading model ...") + else: + # The model is deliberately not loaded without a key. + log(f"FATAL: {state.detail}") + httpd.serve_forever() + + +if __name__ == "__main__": + main() diff --git a/src/vla_sim/hf_cache/.gitkeep b/src/vla_sim/hf_cache/.gitkeep new file mode 100644 index 000000000..e69de29bb diff --git a/src/vla_sim/launch/runtime.launch.xml b/src/vla_sim/launch/runtime.launch.xml new file mode 100644 index 000000000..93d6eacc3 --- /dev/null +++ b/src/vla_sim/launch/runtime.launch.xml @@ -0,0 +1,27 @@ + + + + + + + + + diff --git a/src/vla_sim/models/.gitkeep b/src/vla_sim/models/.gitkeep new file mode 100644 index 000000000..e69de29bb diff --git a/src/vla_sim/objectives/stack_cubes_with_the_vla_policy.xml b/src/vla_sim/objectives/stack_cubes_with_the_vla_policy.xml new file mode 100644 index 000000000..0ab400ee9 --- /dev/null +++ b/src/vla_sim/objectives/stack_cubes_with_the_vla_policy.xml @@ -0,0 +1,56 @@ + + + + + + + + + + + + + + + + + + + diff --git a/src/vla_sim/package.xml b/src/vla_sim/package.xml new file mode 100644 index 000000000..36c2ebfc5 --- /dev/null +++ b/src/vla_sim/package.xml @@ -0,0 +1,46 @@ + + + vla_sim + 9.5.0 + + + MuJoCo simulation configuration package for a Kinova Gen3 (7-DoF) with a + Robotiq 2F-85, stacking colored cubes on command using a + vision-language-action policy. + + + MoveIt Pro Maintainer + + BSD-3-Clause + + ament_cmake + + joint_trajectory_admittance_controller + kortex_description + moveit_pro_behavior + moveit_studio_agent + picknik_mujoco_ros + robotiq_description + velocity_force_controller + + moveit_pro_ml_msgs + python3-numpy + python3-opencv + python3-requests + rclpy + sensor_msgs + std_msgs + trajectory_msgs + + ament_lint_auto + + ament_cmake_copyright + ament_cmake_lint_cmake + ament_cmake_pytest + ament_flake8 + picknik_ament_copyright + + + ament_cmake + + diff --git a/src/vla_sim/script/get_action_chunk_adapter.py b/src/vla_sim/script/get_action_chunk_adapter.py new file mode 100755 index 000000000..c0192b98f --- /dev/null +++ b/src/vla_sim/script/get_action_chunk_adapter.py @@ -0,0 +1,429 @@ +#!/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. + +"""GetActionChunk adapter: bridges ExecutePolicy to the inference server. + +Serves moveit_pro_ml_msgs/srv/GetActionChunk inside the MoveIt Pro agent +container and forwards each request over HTTP to the `inference_server` +container (docker/vla_inference_server.py), which owns torch and the +checkpoint. Requests carry the deployment's shared MOVEIT_FRONTEND_KEY as a +bearer token, the contract the server enforces on /infer. Policy-agnostic and +lightweight: no ML dependencies, so it always runs with the config. + +A failed request answers with status ERROR, which fails the run, and the +response message is what ExecutePolicy shows the operator as the reason. The +messages therefore distinguish "server not running" from "model still +loading" from a server-reported inference error. +""" + +import base64 +import ipaddress +import os +from urllib.parse import urlsplit, urlunsplit + +import cv2 +import numpy as np +import requests + +import rclpy +from rclpy.executors import ExternalShutdownException +from rclpy.node import Node +from std_msgs.msg import Float64MultiArray, MultiArrayDimension +from trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint + +from moveit_pro_ml_msgs.srv import GetActionChunk + +DEFAULT_INFER_URL = "http://127.0.0.1:8973/infer" + +# Bound on server-supplied text relayed into the UI-bound response message. +MAX_SERVER_DETAIL_CHARS = 2000 + + +def clip_detail(text: str) -> str: + """Clip server-supplied text so a runaway error body cannot flood the UI.""" + if len(text) <= MAX_SERVER_DETAIL_CHARS: + return text + return text[:MAX_SERVER_DETAIL_CHARS] + " [...]" + + +def resolve_http_timeout(value: float) -> float: + """Reject a non-positive or non-finite http_timeout at startup, not per call.""" + if not np.isfinite(value) or value <= 0.0: + raise ValueError( + f"the http_timeout parameter must be a positive number of " + f"seconds, got {value!r}" + ) + return float(value) + + +def resolve_infer_url(value: str) -> str: + """Require infer_url to address a server on this machine. + + Requests to it carry the deployment's shared key, and the server speaks + plain HTTP, so a host off this machine would put that key on the wire in + cleartext. Every supported layout is still reachable: the agent container + runs with host networking, so loopback here reaches the product sidecar, + another container publishing to loopback, or a bare host process. + + Only a literal loopback address is accepted. A name resolves at connect + time, not here, so a check on `localhost` would be a check on whatever the + resolver returns later. Only http is accepted too: the server listens + without TLS, so an https target could not complete a handshake with it. + + @return: The address to call, rebuilt from the parts that were checked. + """ + parts = urlsplit(value) + if parts.scheme != "http": + raise ValueError( + f"the infer_url parameter must be an http URL: the inference server " + f"listens without TLS, and the target stays on this machine, so there " + f"is no network hop to encrypt. Got {value!r}" + ) + # `urlsplit` and the HTTP client disagree about where the authority ends: + # `urlsplit` reads `evil.example\@127.0.0.1` as credentials followed by a + # loopback host, while the client stops at the backslash and connects to + # evil.example. Credentials have no use here, since the key travels in a + # header, so rejecting both characters leaves them nothing to disagree on. + if "@" in parts.netloc or "\\" in parts.netloc: + raise ValueError( + f"the infer_url parameter must be a plain scheme://host:port URL " + f"carrying no credentials, got {value!r}" + ) + try: + host = ipaddress.ip_address(parts.hostname or "") + except ValueError: + raise ValueError( + f"the infer_url parameter must address a loopback IP literal such as " + f"{DEFAULT_INFER_URL}, got {value!r}" + ) from None + if not host.is_loopback: + raise ValueError( + f"the infer_url parameter must stay on this machine: {host} is not a " + f"loopback address. Serving the policy from another host is not " + f"supported" + ) + try: + port = parts.port + except ValueError: + raise ValueError( + f"the infer_url parameter has an invalid port, got {value!r}" + ) from None + # Rebuilt rather than returned as given: `urlsplit` drops tab and newline + # characters before parsing, so the original string can still carry bytes + # the client would read differently. Any fragment is dropped with them, + # since it is never sent to the server. + literal = f"[{host}]" if host.version == 6 else str(host) + authority = literal if port is None else f"{literal}:{port}" + return urlunsplit((parts.scheme, authority, parts.path, parts.query, "")) + + +class RequestError(Exception): + """Operator-facing failure; the message becomes the service response message.""" + + +def encode_jpeg_b64(img) -> str: + """sensor_msgs/Image -> base64 JPEG. Accepts rgb8/bgr8/rgba8/bgra8 frames.""" + encoding = str(img.encoding).lower() + if encoding not in ("rgb8", "bgr8", "rgba8", "bgra8"): + raise ValueError(f"unsupported image encoding '{img.encoding}'") + channels = 4 if encoding in ("rgba8", "bgra8") else 3 + # step is the row stride in bytes; slice off any row padding before reshaping. + rows = np.frombuffer(bytes(img.data), dtype=np.uint8).reshape(img.height, img.step) + arr = rows[:, : img.width * channels].reshape(img.height, img.width, channels)[ + :, :, :3 + ] + # cv2.imencode expects BGR input, so bgr frames pass through and rgb frames flip once. + if encoding.startswith("rgb"): + arr = arr[:, :, ::-1] + ok, buf = cv2.imencode(".jpg", np.ascontiguousarray(arr)) + if not ok: + raise RuntimeError("cv2.imencode failed") + return base64.b64encode(buf.tobytes()).decode("ascii") + + +class GetActionChunkAdapter(Node): + def __init__(self) -> None: + super().__init__("get_action_chunk_adapter") + self.infer_url = resolve_infer_url( + self.declare_parameter("infer_url", DEFAULT_INFER_URL).value + ) + # Total HTTP budget, split into connect + read at the call site. Keep + # it strictly below the caller's service timeout (ExecutePolicy's + # policy_call_timeout, 10.0 in stack_cubes_with_the_vla_policy.xml): a hung server + # must not wedge this single-threaded node past the point the caller + # has already given up, or the next run's first request queues behind + # the stale call. + self.http_timeout = resolve_http_timeout( + self.declare_parameter("http_timeout", 9.0).value + ) + # The server authenticates /infer with the deployment's shared key + # (same contract as the web backend endpoints). The agent container + # always carries it; when it is absent the server's 401 detail flows + # into the objective's on-screen message, so no local check is needed. + self._frontend_key = os.environ.get("MOVEIT_FRONTEND_KEY", "").strip() + service_name = self.declare_parameter("service_name", "/get_action_chunk").value + self.create_service(GetActionChunk, service_name, self._on_request) + self._calls = 0 + self.get_logger().info(f"serving '{service_name}' -> {self.infer_url}") + + def _on_request(self, request, response): + try: + self._fill_response(request, response) + except RequestError as exc: + response.status = GetActionChunk.Response.ERROR + response.message = str(exc) + self.get_logger().error(response.message) + except Exception as exc: + response.status = GetActionChunk.Response.ERROR + response.message = f"adapter failed: {type(exc).__name__}: {exc}" + self.get_logger().error(response.message) + return response + + def _build_payload(self, request) -> dict: + """Turn the service request into the server's JSON payload.""" + if len(request.images) != len(request.image_names): + raise RequestError( + f"images ({len(request.images)}) and image_names " + f"({len(request.image_names)}) length mismatch" + ) + # The payload state is built from positions while the returned chunk is + # validated and labeled with the joint names, so the two must agree. + if len(request.robot_state.name) != len(request.robot_state.position): + raise RequestError( + f"robot_state carries {len(request.robot_state.position)} positions " + f"but {len(request.robot_state.name)} joint names" + ) + try: + images = { + request.image_names[i]: encode_jpeg_b64(img) + for i, img in enumerate(request.images) + } + except (ValueError, RuntimeError) as exc: + raise RequestError(f"image encode failed: {exc}") from exc + payload = { + # The request carries the policy's full trained state: the arm group's + # joints plus the gripper joint appended last when the Objective configures one. + "state": list(request.robot_state.position), + "task": request.prompt, + "images": images, + "new_episode": bool(request.new_episode), + } + + # RTC carryover: forward the previous chunk's unexecuted tail and the overlap depth. + # A Float64MultiArray is a flat row-major buffer plus a layout, so reshape it back + # into (steps, action width) rows; it arrives empty on the first call and when RTC is off. + # previous_anchor_state is intentionally not forwarded: the carryover stays in the + # policy's own action space, where lerobot's RTC guidance operates without re-anchoring. + prev = request.previous_action_chunk + if prev.data: + if len(prev.layout.dim) != 2: + raise RequestError( + f"previous_action_chunk carries {len(prev.data)} values but " + f"its layout declares {len(prev.layout.dim)} dimensions " + "instead of the required 2 (steps, action width)" + ) + steps, width = prev.layout.dim[0].size, prev.layout.dim[1].size + if steps * width != len(prev.data): + raise RequestError( + f"previous_action_chunk carries {len(prev.data)} values but " + f"its layout declares {steps}x{width} = {steps * width}" + ) + payload["prev_chunk_left_over"] = ( + np.asarray(prev.data, dtype=float).reshape(steps, width).tolist() + ) + payload["inference_delay"] = int(request.frozen_prefix_steps) + # A non-zero guidance_horizon is the Objective overriding the server's + # soft-guidance width; the server maps it onto lerobot's RTC horizon. + if request.guidance_horizon > 0: + payload["guidance_horizon"] = int(request.guidance_horizon) + return payload + + def _post_infer(self, payload: dict) -> dict: + """POST to the inference server; return the parsed response body.""" + # A local container either accepts immediately or is down, so connect + # gets a small slice and the read keeps the rest. The read element is + # a between-bytes timeout, so the split bounds the hung-server case, + # not a slowly trickling body. + connect_s = min(3.0, self.http_timeout / 3.0) + headers = ( + {"Authorization": f"Bearer {self._frontend_key}"} + if self._frontend_key + else {} + ) + try: + resp = requests.post( + self.infer_url, + json=payload, + headers=headers, + timeout=(connect_s, self.http_timeout - connect_s), + ) + except requests.ConnectionError as exc: + raise RequestError( + f"/infer request failed: the inference server at {self.infer_url} " + f"is not reachable ({exc}). Was MoveIt Pro started with " + "--with-inference-server (or the server started with " + "--only-inference-server)? Check with " + "'docker ps --filter name=inference_server'." + ) from exc + except requests.RequestException as exc: + raise RequestError( + f"/infer request failed: {clip_detail(str(exc))}" + ) from exc + # Parse the body before checking the status code: the server reports + # problems as {"error": ...} bodies (load/inference failures with 5xx, + # request-shape rejections with 4xx), and that detail (e.g. "still + # loading the model") is the message the operator needs to see. + try: + data = resp.json() + except ValueError: + data = None + if isinstance(data, dict) and data.get("error"): + raise RequestError(f"/infer error: {clip_detail(str(data['error']))}") + if not resp.ok or not isinstance(data, dict): + raise RequestError(f"/infer request failed: HTTP {resp.status_code}") + return data + + @staticmethod + def _validate_chunk(data: dict, expected_dims: int) -> list: + """Check the returned chunk's shape, values, and dt; return the chunk.""" + chunk = data.get("action_chunk") + if not chunk or "dt" not in data: + raise RequestError( + "/infer response is missing a non-empty action_chunk or dt" + ) + if not isinstance(chunk, list) or not all( + isinstance(step, list) for step in chunk + ): + raise RequestError( + "/infer response's action_chunk is not a steps x dims array" + ) + dt = data["dt"] + if not isinstance(dt, (int, float)) or not np.isfinite(dt) or dt <= 0.0: + raise RequestError( + f"/infer response carries an invalid dt ({dt!r}); playback " + "pacing needs a finite value > 0" + ) + mismatched_width = next( + (len(step) for step in chunk if len(step) != expected_dims), None + ) + if mismatched_width is not None: + raise RequestError( + f"/infer chunk width {mismatched_width} does not match the observed " + f"joint count {expected_dims}" + ) + try: + chunk_arr = np.asarray(chunk, dtype=float) + except (TypeError, ValueError) as exc: + raise RequestError( + f"/infer chunk carries non-numeric action values ({exc})" + ) from exc + if not np.isfinite(chunk_arr).all(): + raise RequestError("/infer chunk carries non-finite action values") + return chunk + + def _fill_response(self, request, response) -> None: + payload = self._build_payload(request) + data = self._post_infer(payload) + chunk = self._validate_chunk(data, len(request.robot_state.name)) + + # The chunk: absolute joint positions. The action columns line up with the request's + # state entries, so the request's joint names are also the chunk's joint names. + traj = JointTrajectory() + traj.joint_names = list(request.robot_state.name) + for step in chunk: + point = JointTrajectoryPoint() + point.positions = [float(v) for v in step] + traj.points.append(point) + response.chunk = traj + response.native_control_period = float(data["dt"]) + + # RTC echo: the normalized model output, one row per step, one column per action + # dimension; the caller returns its unexecuted tail as the next previous_action_chunk. + # An absent or empty echo means the policy offers no RTC carryover, so + # the next request simply arrives without one. + raw = data.get("action_chunk_raw") + if raw: + if ( + not isinstance(raw, list) + or not all( + isinstance(row, list) and len(row) == len(raw[0]) for row in raw + ) + or not raw[0] + ): + raise RequestError( + "/infer response's action_chunk_raw is not a non-empty " + "rectangular steps x dims array" + ) + try: + raw_arr = np.asarray(raw, dtype=float) + except (TypeError, ValueError) as exc: + raise RequestError( + f"/infer response's action_chunk_raw carries non-numeric " + f"values ({exc})" + ) from exc + if not np.isfinite(raw_arr).all(): + raise RequestError( + "/infer response's action_chunk_raw carries non-finite values" + ) + arr = Float64MultiArray() + steps, width = raw_arr.shape + arr.layout.data_offset = 0 + arr.layout.dim = [ + MultiArrayDimension(label="steps", size=steps, stride=steps * width), + MultiArrayDimension(label="dims", size=width, stride=width), + ] + arr.data = raw_arr.ravel().tolist() + response.policy_action_chunk = arr + + # Set last: a malformed raw echo above must fail the whole request, not + # report a produced chunk with the echo missing. + response.status = GetActionChunk.Response.CHUNK_PRODUCED + self._calls += 1 + if self._calls % 10 == 0: + self.get_logger().info(f"get_action_chunk: served {self._calls} chunks") + + +def main() -> None: + rclpy.init() + node = GetActionChunkAdapter() + try: + rclpy.spin(node) + # rclpy.init() installs the signal handlers, so a stack shutdown arrives as + # ExternalShutdownException rather than KeyboardInterrupt, and it has + # already torn the context down: try_shutdown is the idempotent form. + except (KeyboardInterrupt, ExternalShutdownException): + pass + finally: + node.destroy_node() + rclpy.try_shutdown() + + +if __name__ == "__main__": + main() diff --git a/src/vla_sim/test/test_get_action_chunk_adapter.py b/src/vla_sim/test/test_get_action_chunk_adapter.py new file mode 100644 index 000000000..3e1cc032f --- /dev/null +++ b/src/vla_sim/test/test_get_action_chunk_adapter.py @@ -0,0 +1,785 @@ +#!/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. + +"""Tests for the GetActionChunk adapter: image encoding and request/response mapping.""" + +import importlib.util +import os +import sys +import unittest +from pathlib import Path +from unittest.mock import MagicMock, patch + +import rclpy +import requests +from sensor_msgs.msg import Image, JointState +from std_msgs.msg import Float64MultiArray, MultiArrayDimension, MultiArrayLayout + +from moveit_pro_ml_msgs.srv import GetActionChunk + +# get_action_chunk_adapter.py is a standalone ROS executable (install(PROGRAMS ...) in +# CMakeLists.txt), not part of an importable Python package, so it's loaded by file path. +_SCRIPT_PATH = ( + Path(__file__).resolve().parent.parent / "script" / "get_action_chunk_adapter.py" +) +_spec = importlib.util.spec_from_file_location("get_action_chunk_adapter", _SCRIPT_PATH) +get_action_chunk_adapter = importlib.util.module_from_spec(_spec) +sys.modules["get_action_chunk_adapter"] = get_action_chunk_adapter +_spec.loader.exec_module(get_action_chunk_adapter) + +GetActionChunkAdapter = get_action_chunk_adapter.GetActionChunkAdapter +encode_jpeg_b64 = get_action_chunk_adapter.encode_jpeg_b64 +resolve_http_timeout = get_action_chunk_adapter.resolve_http_timeout +resolve_infer_url = get_action_chunk_adapter.resolve_infer_url + + +def make_image(encoding: str, height: int = 2, width: int = 2) -> Image: + """A flat-colored sensor_msgs/Image with no row padding (step == width * channels).""" + channels = 4 if encoding in ("rgba8", "bgra8") else 3 + img = Image() + img.encoding = encoding + img.height = height + img.width = width + img.step = width * channels + img.data = bytes([128] * (height * img.step)) + return img + + +def make_request(**overrides) -> GetActionChunk.Request: + request = GetActionChunk.Request() + request.robot_state = JointState(name=["j1", "j2"], position=[0.1, 0.2]) + request.images = [make_image("rgb8")] + request.image_names = ["front"] + request.prompt = "stack the blocks" + request.new_episode = False + for key, value in overrides.items(): + setattr(request, key, value) + return request + + +class TestEncodeJpegB64(unittest.TestCase): + """encode_jpeg_b64: sensor_msgs/Image -> base64 JPEG round trip.""" + + def test_unsupported_encoding_raises(self) -> None: + """A mono8 (or any non-{r,b}gb[a]8) image is rejected, not silently reinterpreted.""" + img = make_image("mono8") + with self.assertRaises(ValueError): + encode_jpeg_b64(img) + + def test_rgb8_and_bgr8_roundtrip_to_same_pixels(self) -> None: + """rgb8 and bgr8 inputs carrying the same visual color decode to matching JPEG bytes.""" + rgb = make_image("rgb8") + rgb.data = bytes([10, 20, 30] * (rgb.height * rgb.width)) + bgr = make_image("bgr8") + bgr.data = bytes([30, 20, 10] * (bgr.height * bgr.width)) + + rgb_b64 = encode_jpeg_b64(rgb) + bgr_b64 = encode_jpeg_b64(bgr) + + self.assertEqual(rgb_b64, bgr_b64) + + def test_rgba8_drops_alpha_channel(self) -> None: + """Encoding a 4-channel frame does not crash and yields a valid JPEG (3-channel).""" + img = make_image("rgba8") + img.data = bytes([10, 20, 30, 255] * (img.height * img.width)) + b64 = encode_jpeg_b64(img) + self.assertTrue(len(b64) > 0) + + def test_row_padding_is_stripped(self) -> None: + """step wider than width*channels (row padding) must not corrupt the decoded pixels.""" + img = make_image("rgb8") + pad = 4 + img.step = img.width * 3 + pad + img.data = bytes([10, 20, 30] * img.width + [0] * pad) * img.height + padded_b64 = encode_jpeg_b64(img) + + unpadded = make_image("rgb8") + unpadded.data = bytes([10, 20, 30] * (unpadded.height * unpadded.width)) + unpadded_b64 = encode_jpeg_b64(unpadded) + + self.assertEqual(padded_b64, unpadded_b64) + + +class TestResolveHttpTimeout(unittest.TestCase): + """resolve_http_timeout: unusable values fail at startup, naming the parameter.""" + + def test_positive_value_passes_through(self) -> None: + """A positive timeout is returned unchanged.""" + self.assertEqual(resolve_http_timeout(9.0), 9.0) + + def test_zero_raises_naming_the_parameter(self) -> None: + """A zero timeout fails at construction, not as a per-call generic error.""" + with self.assertRaises(ValueError) as ctx: + resolve_http_timeout(0.0) + self.assertIn("http_timeout", str(ctx.exception)) + + +class TestResolveInferUrl(unittest.TestCase): + """resolve_infer_url: requests carry the deployment key over plain HTTP, so + the target has to stay on this machine.""" + + def test_default_url_is_accepted(self) -> None: + """The shipped default must pass its own check.""" + self.assertEqual( + resolve_infer_url(get_action_chunk_adapter.DEFAULT_INFER_URL), + get_action_chunk_adapter.DEFAULT_INFER_URL, + ) + + def test_ipv6_loopback_is_accepted(self) -> None: + """A server bound to ::1 is just as local as one on 127.0.0.1.""" + url = "http://[::1]:8973/infer" + self.assertEqual(resolve_infer_url(url), url) + + def test_loopback_range_beyond_the_first_address_is_accepted(self) -> None: + """The whole 127.0.0.0/8 range is loopback, not just 127.0.0.1.""" + url = "http://127.0.0.5:8973/infer" + self.assertEqual(resolve_infer_url(url), url) + + def test_remote_host_is_rejected(self) -> None: + """The shared key would go out in cleartext to a machine we do not own.""" + with self.assertRaises(ValueError) as ctx: + resolve_infer_url("http://10.0.0.7:8973/infer") + self.assertIn("loopback", str(ctx.exception)) + + def test_remote_https_host_is_rejected(self) -> None: + """TLS is not the bar: serving off this machine is not supported.""" + with self.assertRaises(ValueError): + resolve_infer_url("https://10.0.0.7:8973/infer") + + def test_https_is_rejected_even_for_loopback(self) -> None: + """The server listens without TLS, so an https target could only ever + fail the handshake. Rejecting it here names the scheme as the problem, + rather than surfacing later as an unreachable server.""" + for url in ("https://127.0.0.1:8973/infer", "https://[::1]:8973/infer"): + with self.assertRaises(ValueError, msg=url) as ctx: + resolve_infer_url(url) + self.assertIn("http URL", str(ctx.exception), msg=url) + + def test_hostname_is_rejected_even_when_it_names_loopback(self) -> None: + """'localhost' resolves at connect time, so accepting it would accept + whatever the resolver returns then, not what is checked here.""" + with self.assertRaises(ValueError) as ctx: + resolve_infer_url("http://localhost:8973/infer") + self.assertIn("IP literal", str(ctx.exception)) + + def test_userinfo_cannot_disguise_a_remote_host(self) -> None: + """Everything before '@' is userinfo: this URL addresses evil.com.""" + with self.assertRaises(ValueError) as ctx: + resolve_infer_url("http://127.0.0.1@evil.com/infer") + self.assertIn("no credentials", str(ctx.exception)) + + def test_backslash_in_the_authority_is_rejected(self) -> None: + """urlsplit reads the backslash as an ordinary userinfo character and + reports 127.0.0.1, while the HTTP client ends the authority there and + connects to evil.example. Accepting this sends the deployment key to + that host in cleartext.""" + for url in ( + "http://evil.example\\@127.0.0.1:8973/infer", + "http://evil.example\\@[::1]:8973/infer", + ): + with self.assertRaises(ValueError, msg=url) as ctx: + resolve_infer_url(url) + self.assertIn("no credentials", str(ctx.exception), msg=url) + + def test_the_returned_url_is_rebuilt_from_the_checked_parts(self) -> None: + """urlsplit drops tab and newline characters before it parses, so a URL + handed back as given can still carry bytes that were never checked.""" + self.assertEqual( + resolve_infer_url("http://127.0.0.1:8973/infer\n"), + "http://127.0.0.1:8973/infer", + ) + + def test_a_port_the_client_cannot_use_is_rejected_here(self) -> None: + """Out-of-range and non-numeric ports otherwise surface much later, as + a per-request client error that does not name the parameter at fault.""" + for url in ("http://127.0.0.1:99999/infer", "http://127.0.0.1:8973-8975/infer"): + with self.assertRaises(ValueError, msg=url) as ctx: + resolve_infer_url(url) + self.assertIn("invalid port", str(ctx.exception), msg=url) + + def test_ipv4_mapped_ipv6_cannot_disguise_a_remote_host(self) -> None: + """::ffff:10.0.0.7 parses as an IP literal but is not loopback.""" + with self.assertRaises(ValueError) as ctx: + resolve_infer_url("http://[::ffff:10.0.0.7]:8973/infer") + self.assertIn("loopback", str(ctx.exception)) + + def test_alternate_spellings_of_loopback_are_rejected(self) -> None: + """Decimal and hex forms of 127.0.0.1 are what a bypass looks like; + requiring dotted-quad keeps the accepted set to one spelling.""" + for url in ("http://2130706433:8973/infer", "http://0x7f000001:8973/infer"): + with self.assertRaises(ValueError, msg=url): + resolve_infer_url(url) + + def test_non_http_scheme_is_rejected(self) -> None: + """The adapter POSTs with requests; a file:// target is not a server.""" + with self.assertRaises(ValueError) as ctx: + resolve_infer_url("file:///etc/passwd") + self.assertIn("http URL", str(ctx.exception)) + + +class TestOnRequest(unittest.TestCase): + """GetActionChunkAdapter._on_request: HTTP call shaping and response translation.""" + + @classmethod + def setUpClass(cls) -> None: + rclpy.init() + + @classmethod + def tearDownClass(cls) -> None: + rclpy.shutdown() + + def setUp(self) -> None: + self.node = GetActionChunkAdapter() + + def tearDown(self) -> None: + self.node.destroy_node() + + @patch("get_action_chunk_adapter.requests.post") + def test_successful_chunk_populates_trajectory(self, mock_post: MagicMock) -> None: + """A valid /infer response becomes a JointTrajectory with matching joint names.""" + mock_post.return_value.json.return_value = { + "action_chunk": [[0.5, 0.6], [0.7, 0.8]], + "dt": 0.05, + } + request = make_request() + response = self.node._on_request(request, GetActionChunk.Response()) + + self.assertEqual( + response.status, GetActionChunk.Response.CHUNK_PRODUCED, response.message + ) + self.assertEqual(list(response.chunk.joint_names), ["j1", "j2"]) + self.assertEqual(len(response.chunk.points), 2) + self.assertEqual(list(response.chunk.points[0].positions), [0.5, 0.6]) + self.assertAlmostEqual(response.native_control_period, 0.05) + + @patch("get_action_chunk_adapter.requests.post") + def test_prompt_forwarded_as_task_field(self, mock_post: MagicMock) -> None: + """The fixed request.prompt field is sent to /infer under the 'task' key.""" + mock_post.return_value.json.return_value = { + "action_chunk": [[0.0, 0.0]], + "dt": 0.05, + } + request = make_request(prompt="pick the red cube") + self.node._on_request(request, GetActionChunk.Response()) + + sent_payload = mock_post.call_args.kwargs["json"] + self.assertEqual(sent_payload["task"], "pick the red cube") + + @patch("get_action_chunk_adapter.requests.post") + def test_observation_forwarded_in_payload(self, mock_post: MagicMock) -> None: + """The joint positions, image names, and new_episode flag all reach the + payload, posted to the node's configured infer_url.""" + mock_post.return_value.json.return_value = { + "action_chunk": [[0.0, 0.0]], + "dt": 0.05, + } + request = make_request(new_episode=True) + self.node._on_request(request, GetActionChunk.Response()) + + (url,) = mock_post.call_args.args + self.assertEqual(url, self.node.infer_url) + sent_payload = mock_post.call_args.kwargs["json"] + self.assertEqual(sent_payload["state"], [0.1, 0.2]) + self.assertEqual(list(sent_payload["images"].keys()), ["front"]) + self.assertTrue(sent_payload["new_episode"]) + + def test_unsupported_image_encoding_fails_without_http_call(self) -> None: + """An image the encoder can't handle fails locally; /infer is never called.""" + request = make_request(images=[make_image("mono8")]) + with patch("get_action_chunk_adapter.requests.post") as mock_post: + response = self.node._on_request(request, GetActionChunk.Response()) + mock_post.assert_not_called() + self.assertEqual(response.status, GetActionChunk.Response.ERROR) + self.assertIn("image encode failed", response.message) + + def test_images_and_image_names_length_mismatch_fails_without_http_call( + self, + ) -> None: + """A malformed request (arrays not lined up by index, per the .srv contract) + fails locally instead of raising IndexError out of the service callback.""" + request = make_request(images=[make_image("rgb8"), make_image("rgb8")]) + with patch("get_action_chunk_adapter.requests.post") as mock_post: + response = self.node._on_request(request, GetActionChunk.Response()) + mock_post.assert_not_called() + self.assertEqual(response.status, GetActionChunk.Response.ERROR) + self.assertIn("length mismatch", response.message) + + def test_state_name_position_length_mismatch_fails_without_http_call(self) -> None: + """Mismatched robot_state arrays fail locally naming the real defect; + the payload state comes from positions while the chunk is validated and + labeled with the joint names.""" + request = make_request( + robot_state=JointState(name=["j1", "j2", "j3"], position=[0.1, 0.2]) + ) + with patch("get_action_chunk_adapter.requests.post") as mock_post: + response = self.node._on_request(request, GetActionChunk.Response()) + mock_post.assert_not_called() + self.assertEqual(response.status, GetActionChunk.Response.ERROR) + self.assertIn("2 positions", response.message) + self.assertIn("3 joint names", response.message) + + @patch("get_action_chunk_adapter.requests.post") + def test_connection_error_names_the_server_and_a_check_command( + self, mock_post: MagicMock + ) -> None: + """A refused connection tells the operator which server is down and how to check it.""" + mock_post.side_effect = requests.ConnectionError("refused") + response = self.node._on_request(make_request(), GetActionChunk.Response()) + + self.assertEqual(response.status, GetActionChunk.Response.ERROR) + self.assertIn("/infer request failed", response.message) + self.assertIn("inference_server", response.message) + self.assertIn("--with-inference-server", response.message) + + @patch("get_action_chunk_adapter.requests.post") + def test_request_carries_frontend_key_bearer_token( + self, mock_post: MagicMock + ) -> None: + """/infer requests present MOVEIT_FRONTEND_KEY as the bearer token the + server's auth check expects.""" + mock_post.return_value.json.return_value = { + "action_chunk": [[0.0, 0.0]], + "dt": 0.05, + } + with patch.dict(os.environ, {"MOVEIT_FRONTEND_KEY": "secret-key"}): + node = GetActionChunkAdapter() + try: + node._on_request(make_request(), GetActionChunk.Response()) + finally: + node.destroy_node() + + self.assertEqual( + mock_post.call_args.kwargs["headers"], + {"Authorization": "Bearer secret-key"}, + ) + + @patch("get_action_chunk_adapter.requests.post") + def test_missing_key_sends_no_authorization_header( + self, mock_post: MagicMock + ) -> None: + """Without a key the request goes out bare; the server's 401 detail then + names the fix through the normal error path.""" + mock_post.return_value.json.return_value = { + "action_chunk": [[0.0, 0.0]], + "dt": 0.05, + } + with patch.dict(os.environ): + os.environ.pop("MOVEIT_FRONTEND_KEY", None) + node = GetActionChunkAdapter() + try: + node._on_request(make_request(), GetActionChunk.Response()) + finally: + node.destroy_node() + + self.assertEqual(mock_post.call_args.kwargs["headers"], {}) + + @patch("get_action_chunk_adapter.requests.post") + def test_timeout_surfaces_as_failure(self, mock_post: MagicMock) -> None: + """A network-level failure other than a refused connection is reported, not raised.""" + mock_post.side_effect = requests.Timeout("timed out") + response = self.node._on_request(make_request(), GetActionChunk.Response()) + + self.assertEqual(response.status, GetActionChunk.Response.ERROR) + self.assertIn("/infer request failed", response.message) + + @patch("get_action_chunk_adapter.requests.post") + def test_http_timeout_tracks_the_parameter(self, mock_post: MagicMock) -> None: + """The connect + read split sums to the http_timeout parameter (9.0 + default, strictly below the caller's 10.0 service timeout), so a hung + server cannot wedge the single-threaded node past the caller.""" + mock_post.return_value.json.return_value = { + "action_chunk": [[0.0, 0.0]], + "dt": 0.05, + } + self.node._on_request(make_request(), GetActionChunk.Response()) + + connect_s, read_s = mock_post.call_args.kwargs["timeout"] + self.assertEqual(connect_s, 3.0) + self.assertEqual(connect_s + read_s, 9.0) + + @patch("get_action_chunk_adapter.requests.post") + def test_server_error_field_surfaces_as_failure(self, mock_post: MagicMock) -> None: + """A server-reported {"error": ...} body fails the request rather than being ignored. + + This is also the path that relays "still loading the model" (503) and + "model load failed" (500) to the operator.""" + mock_post.return_value.json.return_value = {"error": "checkpoint not loaded"} + response = self.node._on_request(make_request(), GetActionChunk.Response()) + + self.assertEqual(response.status, GetActionChunk.Response.ERROR) + self.assertIn("checkpoint not loaded", response.message) + + @patch("get_action_chunk_adapter.requests.post") + def test_huge_server_error_is_truncated(self, mock_post: MagicMock) -> None: + """A runaway server error string is clipped before it reaches the UI.""" + mock_post.return_value.json.return_value = {"error": "x" * 100_000} + response = self.node._on_request(make_request(), GetActionChunk.Response()) + + self.assertEqual(response.status, GetActionChunk.Response.ERROR) + self.assertLess(len(response.message), 3000) + self.assertIn("[...]", response.message) + + @patch("get_action_chunk_adapter.requests.post") + def test_non_json_error_response_reports_http_status( + self, mock_post: MagicMock + ) -> None: + """A non-JSON body on a failed status still produces a diagnosable message.""" + mock_post.return_value.json.side_effect = ValueError("no json") + mock_post.return_value.ok = False + mock_post.return_value.status_code = 502 + response = self.node._on_request(make_request(), GetActionChunk.Response()) + + self.assertEqual(response.status, GetActionChunk.Response.ERROR) + self.assertIn("HTTP 502", response.message) + + @patch("get_action_chunk_adapter.requests.post") + def test_missing_dt_fails(self, mock_post: MagicMock) -> None: + """A response with a chunk but no dt is rejected rather than defaulting silently.""" + mock_post.return_value.json.return_value = {"action_chunk": [[0.0, 0.0]]} + response = self.node._on_request(make_request(), GetActionChunk.Response()) + + self.assertEqual(response.status, GetActionChunk.Response.ERROR) + self.assertIn("missing", response.message) + + @patch("get_action_chunk_adapter.requests.post") + def test_chunk_width_mismatch_fails(self, mock_post: MagicMock) -> None: + """A chunk whose column count doesn't match the requested joint count is rejected.""" + mock_post.return_value.json.return_value = { + "action_chunk": [[0.0, 0.0, 0.0]], # 3 columns, request has 2 joints + "dt": 0.05, + } + response = self.node._on_request(make_request(), GetActionChunk.Response()) + + self.assertEqual(response.status, GetActionChunk.Response.ERROR) + self.assertIn("does not match", response.message) + + @patch("get_action_chunk_adapter.requests.post") + def test_chunk_width_mismatch_reports_the_offending_row( + self, mock_post: MagicMock + ) -> None: + """The error names the row that actually mismatched, not always row 0.""" + mock_post.return_value.json.return_value = { + # row 0 matches the request's 2 joints; row 1 is the actual offender. + "action_chunk": [[0.0, 0.0], [0.0, 0.0, 0.0]], + "dt": 0.05, + } + response = self.node._on_request(make_request(), GetActionChunk.Response()) + + self.assertEqual(response.status, GetActionChunk.Response.ERROR) + self.assertIn("chunk width 3", response.message) + + @patch("get_action_chunk_adapter.requests.post") + def test_scalar_chunk_rows_fail_naming_action_chunk( + self, mock_post: MagicMock + ) -> None: + """Rows that aren't lists (a flat or scalar chunk from a custom server) + are rejected naming action_chunk, not as a generic adapter failure.""" + mock_post.return_value.json.return_value = { + "action_chunk": [0.0, 0.1], + "dt": 0.05, + } + response = self.node._on_request(make_request(), GetActionChunk.Response()) + + self.assertEqual(response.status, GetActionChunk.Response.ERROR) + self.assertIn("action_chunk is not a steps x dims array", response.message) + + @patch("get_action_chunk_adapter.requests.post") + def test_non_numeric_chunk_values_fail_naming_action_chunk( + self, mock_post: MagicMock + ) -> None: + """String action values are rejected naming the chunk, not as a raw + numpy conversion error.""" + mock_post.return_value.json.return_value = { + "action_chunk": [["a", "b"]], + "dt": 0.05, + } + response = self.node._on_request(make_request(), GetActionChunk.Response()) + + self.assertEqual(response.status, GetActionChunk.Response.ERROR) + self.assertIn("non-numeric action values", response.message) + + @patch("get_action_chunk_adapter.requests.post") + def test_non_finite_chunk_fails(self, mock_post: MagicMock) -> None: + """NaN/inf action values are rejected before they reach the controller.""" + mock_post.return_value.json.return_value = { + "action_chunk": [[0.0, float("nan")]], + "dt": 0.05, + } + response = self.node._on_request(make_request(), GetActionChunk.Response()) + + self.assertEqual(response.status, GetActionChunk.Response.ERROR) + self.assertIn("non-finite", response.message) + + @patch("get_action_chunk_adapter.requests.post") + def test_non_positive_dt_fails(self, mock_post: MagicMock) -> None: + """A dt of zero (or below) is rejected: it cannot pace chunk playback.""" + mock_post.return_value.json.return_value = { + "action_chunk": [[0.0, 0.0]], + "dt": 0.0, + } + response = self.node._on_request(make_request(), GetActionChunk.Response()) + + self.assertEqual(response.status, GetActionChunk.Response.ERROR) + self.assertIn("invalid dt", response.message) + + @patch("get_action_chunk_adapter.requests.post") + def test_non_finite_dt_fails(self, mock_post: MagicMock) -> None: + """A NaN dt is rejected before it reaches trajectory timing.""" + mock_post.return_value.json.return_value = { + "action_chunk": [[0.0, 0.0]], + "dt": float("nan"), + } + response = self.node._on_request(make_request(), GetActionChunk.Response()) + + self.assertEqual(response.status, GetActionChunk.Response.ERROR) + self.assertIn("invalid dt", response.message) + + @patch("get_action_chunk_adapter.requests.post") + def test_previous_chunk_reshaped_and_forwarded(self, mock_post: MagicMock) -> None: + """A populated previous_action_chunk is reshaped from its flat layout before sending.""" + mock_post.return_value.json.return_value = { + "action_chunk": [[0.0, 0.0]], + "dt": 0.05, + } + prev = Float64MultiArray( + layout=MultiArrayLayout( + dim=[ + MultiArrayDimension(label="steps", size=2, stride=4), + MultiArrayDimension(label="dims", size=2, stride=2), + ] + ), + data=[1.0, 2.0, 3.0, 4.0], + ) + request = make_request(previous_action_chunk=prev, frozen_prefix_steps=3) + self.node._on_request(request, GetActionChunk.Response()) + + sent_payload = mock_post.call_args.kwargs["json"] + self.assertEqual(sent_payload["prev_chunk_left_over"], [[1.0, 2.0], [3.0, 4.0]]) + self.assertEqual(sent_payload["inference_delay"], 3) + + def test_previous_chunk_with_malformed_layout_fails_without_http_call(self) -> None: + """A populated previous_action_chunk whose layout is not the contract's two + dimensions fails loudly instead of silently dropping the RTC carryover.""" + prev = Float64MultiArray( + layout=MultiArrayLayout( + dim=[MultiArrayDimension(label="flat", size=4, stride=4)] + ), + data=[1.0, 2.0, 3.0, 4.0], + ) + request = make_request(previous_action_chunk=prev, frozen_prefix_steps=3) + with patch("get_action_chunk_adapter.requests.post") as mock_post: + response = self.node._on_request(request, GetActionChunk.Response()) + mock_post.assert_not_called() + + self.assertEqual(response.status, GetActionChunk.Response.ERROR) + self.assertIn("1 dimensions", response.message) + self.assertIn("2 (steps, action width)", response.message) + + def test_previous_chunk_with_wrong_element_count_fails_without_http_call( + self, + ) -> None: + """A previous_action_chunk whose data length contradicts its declared + layout fails with the mismatch named, not a raw reshape error.""" + prev = Float64MultiArray( + layout=MultiArrayLayout( + dim=[ + MultiArrayDimension(label="steps", size=2, stride=6), + MultiArrayDimension(label="dims", size=3, stride=3), + ] + ), + data=[1.0, 2.0, 3.0, 4.0], + ) + request = make_request(previous_action_chunk=prev, frozen_prefix_steps=3) + with patch("get_action_chunk_adapter.requests.post") as mock_post: + response = self.node._on_request(request, GetActionChunk.Response()) + mock_post.assert_not_called() + + self.assertEqual(response.status, GetActionChunk.Response.ERROR) + self.assertIn("4 values", response.message) + self.assertIn("2x3", response.message) + + @patch("get_action_chunk_adapter.requests.post") + def test_empty_previous_chunk_omits_rtc_fields(self, mock_post: MagicMock) -> None: + """The first call of an episode (empty previous_action_chunk) sends no RTC carryover.""" + mock_post.return_value.json.return_value = { + "action_chunk": [[0.0, 0.0]], + "dt": 0.05, + } + self.node._on_request(make_request(), GetActionChunk.Response()) + + sent_payload = mock_post.call_args.kwargs["json"] + self.assertNotIn("prev_chunk_left_over", sent_payload) + self.assertNotIn("inference_delay", sent_payload) + + @patch("get_action_chunk_adapter.requests.post") + def test_zero_guidance_horizon_omits_the_field(self, mock_post: MagicMock) -> None: + """guidance_horizon=0 defers to the server's own RTC default, per the .srv contract.""" + mock_post.return_value.json.return_value = { + "action_chunk": [[0.0, 0.0]], + "dt": 0.05, + } + self.node._on_request( + make_request(guidance_horizon=0), GetActionChunk.Response() + ) + + sent_payload = mock_post.call_args.kwargs["json"] + self.assertNotIn("guidance_horizon", sent_payload) + + @patch("get_action_chunk_adapter.requests.post") + def test_nonzero_guidance_horizon_is_forwarded(self, mock_post: MagicMock) -> None: + """A nonzero guidance_horizon from the Objective is forwarded as the soft-guidance + width; the server, which knows the inference delay, maps it onto lerobot's horizon. + """ + mock_post.return_value.json.return_value = { + "action_chunk": [[0.0, 0.0]], + "dt": 0.05, + } + self.node._on_request( + make_request(guidance_horizon=7), GetActionChunk.Response() + ) + + sent_payload = mock_post.call_args.kwargs["json"] + self.assertEqual(sent_payload["guidance_horizon"], 7) + + @patch("get_action_chunk_adapter.requests.post") + def test_action_chunk_raw_echoed_as_policy_action_chunk( + self, mock_post: MagicMock + ) -> None: + """When the server echoes action_chunk_raw, it is reflected back as policy_action_chunk.""" + mock_post.return_value.json.return_value = { + "action_chunk": [[0.5, 0.6], [0.7, 0.8]], + "action_chunk_raw": [[-0.1, 0.1], [0.2, -0.2]], + "dt": 0.05, + } + response = self.node._on_request(make_request(), GetActionChunk.Response()) + + arr = response.policy_action_chunk + self.assertEqual([d.size for d in arr.layout.dim], [2, 2]) + self.assertEqual(list(arr.data), [-0.1, 0.1, 0.2, -0.2]) + + @patch("get_action_chunk_adapter.requests.post") + def test_ragged_action_chunk_raw_fails_the_request( + self, mock_post: MagicMock + ) -> None: + """A malformed RTC echo fails the whole request naming the field, never + success with the echo silently missing.""" + mock_post.return_value.json.return_value = { + "action_chunk": [[0.5, 0.6]], + "action_chunk_raw": [[0.1, 0.2], [0.3]], + "dt": 0.05, + } + response = self.node._on_request(make_request(), GetActionChunk.Response()) + + self.assertEqual(response.status, GetActionChunk.Response.ERROR) + self.assertIn("action_chunk_raw", response.message) + + @patch("get_action_chunk_adapter.requests.post") + def test_zero_width_action_chunk_raw_fails_the_request( + self, mock_post: MagicMock + ) -> None: + """An RTC echo with zero-width rows is rejected, not fed back as a + structurally valid but empty carryover.""" + mock_post.return_value.json.return_value = { + "action_chunk": [[0.5, 0.6]], + "action_chunk_raw": [[], []], + "dt": 0.05, + } + response = self.node._on_request(make_request(), GetActionChunk.Response()) + + self.assertEqual(response.status, GetActionChunk.Response.ERROR) + self.assertIn("action_chunk_raw", response.message) + + @patch("get_action_chunk_adapter.requests.post") + def test_empty_action_chunk_raw_succeeds_with_empty_echo( + self, mock_post: MagicMock + ) -> None: + """A present-but-empty echo means no RTC carryover, same as an absent one.""" + mock_post.return_value.json.return_value = { + "action_chunk": [[0.5, 0.6]], + "action_chunk_raw": [], + "dt": 0.05, + } + response = self.node._on_request(make_request(), GetActionChunk.Response()) + + self.assertEqual( + response.status, GetActionChunk.Response.CHUNK_PRODUCED, response.message + ) + self.assertEqual(list(response.policy_action_chunk.data), []) + + @patch("get_action_chunk_adapter.requests.post") + def test_non_finite_action_chunk_raw_fails_the_request( + self, mock_post: MagicMock + ) -> None: + """A NaN in the RTC echo is rejected here, where the field is named, + not one call later as a confusing previous-chunk error.""" + mock_post.return_value.json.return_value = { + "action_chunk": [[0.5, 0.6]], + "action_chunk_raw": [[0.1, float("nan")]], + "dt": 0.05, + } + response = self.node._on_request(make_request(), GetActionChunk.Response()) + + self.assertEqual(response.status, GetActionChunk.Response.ERROR) + self.assertIn("action_chunk_raw", response.message) + self.assertIn("non-finite", response.message) + + @patch("get_action_chunk_adapter.requests.post") + def test_non_numeric_action_chunk_raw_fails_the_request( + self, mock_post: MagicMock + ) -> None: + """Rectangular but non-numeric echo content is reported against the + field, not as a generic adapter failure.""" + mock_post.return_value.json.return_value = { + "action_chunk": [[0.5, 0.6]], + "action_chunk_raw": [["a", "b"]], + "dt": 0.05, + } + response = self.node._on_request(make_request(), GetActionChunk.Response()) + + self.assertEqual(response.status, GetActionChunk.Response.ERROR) + self.assertIn("action_chunk_raw", response.message) + self.assertIn("non-numeric", response.message) + + @patch("get_action_chunk_adapter.requests.post") + def test_no_action_chunk_raw_leaves_policy_action_chunk_empty( + self, mock_post: MagicMock + ) -> None: + """A policy without RTC support (no action_chunk_raw) leaves policy_action_chunk unset.""" + mock_post.return_value.json.return_value = { + "action_chunk": [[0.5, 0.6]], + "dt": 0.05, + } + response = self.node._on_request(make_request(), GetActionChunk.Response()) + + self.assertEqual(list(response.policy_action_chunk.data), []) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/vla_sim/thumbnail.png b/src/vla_sim/thumbnail.png new file mode 100644 index 000000000..6cc52e524 --- /dev/null +++ b/src/vla_sim/thumbnail.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3ce86f92da190a51b6bfa490a68a11aba024f881ba4eaeb3cbe639c6fe3b72f0 +size 195813 diff --git a/src/vla_sim/waypoints/waypoints.yaml b/src/vla_sim/waypoints/waypoints.yaml new file mode 100644 index 000000000..1f3adea18 --- /dev/null +++ b/src/vla_sim/waypoints/waypoints.yaml @@ -0,0 +1,13 @@ +- name: Home + joint_group_names: + - manipulator + joint_state: + header: + frame_id: world + stamp: + sec: 0 + nanosec: 0 + name: [joint_1, joint_2, joint_3, joint_4, joint_5, joint_6, joint_7] + position: [7.2854474427322627e-05, 0.0060998872530424308, -3.1250605484909082, -2.5051932931762071, -2.7164243118882329e-05, 0.95966324117024648, 1.5699864595074029] + velocity: [] + effort: []