Add vla_sim: prompt-driven cube stacking with a vision-language-action policy

.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

.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.*

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`

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/<checkpoint-dir>. 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

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()

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.

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/)

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"