Load the pi0.5 policy at the size it runs at

src/vla_sim/config/vla_serving.yaml

@@ -38,6 +38,17 @@ fps: 10.0
 # order of magnitude slower on cpu.
 device: auto
 
+# Hold the pi0.5 language backbone and vision tower at eight bits, most of the
+# weights: roughly 3 GiB less GPU memory for a little more time per chunk. false
+# loads every weight at full width, and a policy family other than pi0.5 ignores
+# this without saying so.
+#
+# On the stacking objective it has scored within noise of full width, over too
+# few attempts per arm to resolve a small difference. Torchao is imported only
+# when this is on, so an image built before it was a dependency serves at false
+# and reports an error on /health at true; docker/README.md has the rebuild.
+int8: false
+
 # 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.

src/vla_sim/docker/Dockerfile.vla_inference_server

@@ -15,18 +15,36 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
 # 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.
+# torchao supplies the int8 weights vla_serving.yaml can ask for. Its wheel
+# declares no torch dependency, so pip cannot catch a mismatch; torchao skips
+# loading its compiled kernels with a warning when torch is older than the
+# release expects, which leaves the pairing below this file's to keep.
 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" \
+ARG TORCHAO_VERSION=0.17.0
+# This layer moves gigabytes of CUDA wheels, which pip's defaults are not sized
+# for: a read that stalls past the default timeout fails the layer, and a failed
+# layer restarts from zero rather than from where it stopped. Set on the command
+# rather than with ENV, which would carry build tuning into the served image.
+#
+# The cache mount does what --no-cache-dir was doing: it keeps downloaded wheels
+# out of the image, but outside it rather than deleted, so editing a version
+# above costs only the wheels that changed. It is not part of any layer, so the
+# image is the same size either way, and `docker build --no-cache` still starts
+# from an empty one.
+RUN --mount=type=cache,target=/root/.cache/pip \
+    export PIP_DEFAULT_TIMEOUT=120 PIP_RETRIES=10 \
+    && if [ -n "$TORCH_INDEX" ]; then \
+      pip install --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"
+    && pip install "lerobot[pi,smolvla]==0.6.0" \
+      "torch==$TORCH_VERSION" "torchvision==$TORCHVISION_VERSION" \
+      "torchao==$TORCHAO_VERSION"
 
-COPY vla_inference_server.py /app/vla_inference_server.py
+COPY vla_inference_server.py /app/
 WORKDIR /app
 
 # Non-root default for a bare `docker run` (the compose service overrides the

src/vla_sim/docker/README.md

@@ -28,6 +28,23 @@ 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.
 
+Only a *missing* image is built that way, and `moveit_pro build` skips this
+service because its compose profile is off by default, so nothing rebuilds the
+image when this directory changes. The scripts here are mounted rather than read
+from the image, so one that needs a package the existing image predates fails at
+import. After editing the Dockerfile or its pinned versions, drop the image and
+let the next run build it:
+
+```bash
+moveit_pro down
+docker rmi moveit_pro-inference_server:latest
+```
+
+Compose names the image after its project and this service, so the tag above is
+what the launcher builds; `docker images` confirms it. `moveit_pro down` first
+because Docker refuses to remove an image a container still references, and a
+stopped container counts.
+
 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:

src/vla_sim/docker/test_vla_inference_server.py

@@ -77,6 +77,7 @@
     resolve_fps,
     resolve_rtc_horizon,
     resolve_rtc_schedule,
+    trim_vocabulary_heads,
     watch_runtime_inputs,
 )
 
@@ -438,6 +439,42 @@ def test_non_numeric_yaml_values_park_in_config_error(self) -> None:
         self.assertEqual(args.fps, 0.0)
         self.assertEqual(args.state_dim, 0)
 
+    def test_non_boolean_int8_parks_in_config_error(self) -> None:
+        """A quoted or misspelled int8 lands in config_error with the flag off.
+        Coercing it with bool() would read every non-empty string as true, so the
+        server would quantize the policy the operator asked to leave alone."""
+        # GIVEN a serving config whose int8 is a string rather than a boolean
+        with tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False) as f:
+            f.write('int8: "false"\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 the value is reported and the flag stays off
+        self.assertIn("int8", args.config_error)
+        self.assertFalse(args.int8)
+
+    def test_boolean_int8_is_honored(self) -> None:
+        """A real YAML boolean reaches the flag, so the knob works as documented."""
+        # GIVEN a serving config asking for int8
+        with tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False) as f:
+            f.write("int8: true\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 the flag is on and nothing is reported
+        self.assertTrue(args.int8)
+        self.assertEqual(args.config_error, "")
+
     def test_trainer_handoff_fields_are_loaded_from_yaml(self) -> None:
         """The server carries Trainer's exact revision and token requirement."""
         revision = "a" * 40
@@ -574,6 +611,7 @@ def test_local_checkpoint_rejects_revision_before_metadata_resolution(self) -> N
                 guidance_horizon=8,
                 rtc_schedule="EXP",
                 state_dim=8,
+                int8=False,
             )
             with patch("vla_inference_server.PolicyRunner") as runner:
                 load_policy(state, args)
@@ -599,6 +637,7 @@ def _load(self, fps: float, warmup) -> ServerState:
             guidance_horizon=8,
             rtc_schedule="EXP",
             state_dim=8,
+            int8=False,
         )
         runner = MagicMock()
         runner.policy.config.input_features = {}
@@ -666,8 +705,13 @@ def test_policy_and_processors_load_the_pinned_snapshot(self) -> None:
         pre = SimpleNamespace(steps=[])
         post = SimpleNamespace(steps=[])
 
+        policy_config = SimpleNamespace(device="cpu")
         with (
             patch("vla_inference_server.get_policy_class", return_value=policy_class),
+            patch(
+                "vla_inference_server.PreTrainedConfig.from_pretrained",
+                return_value=policy_config,
+            ) as read_config,
             patch(
                 "vla_inference_server.snapshot_download", return_value=snapshot
             ) as download,
@@ -687,7 +731,10 @@ def test_policy_and_processors_load_the_pinned_snapshot(self) -> None:
             )
 
         download.assert_called_once_with("acme/model", revision=revision)
-        policy_class.from_pretrained.assert_called_once_with(snapshot)
+        read_config.assert_called_once_with(snapshot)
+        policy_class.from_pretrained.assert_called_once_with(
+            snapshot, config=policy_config
+        )
         processors.assert_called_once_with(
             policy_cfg=config,
             pretrained_path=snapshot,
@@ -701,8 +748,13 @@ def test_policy_runner_legacy_call_uses_existing_mutable_defaults(self) -> None:
         policy = MagicMock(config=config)
         policy_class = SimpleNamespace(from_pretrained=MagicMock(return_value=policy))
 
+        policy_config = SimpleNamespace(device="cpu")
         with (
             patch("vla_inference_server.get_policy_class", return_value=policy_class),
+            patch(
+                "vla_inference_server.PreTrainedConfig.from_pretrained",
+                return_value=policy_config,
+            ) as read_config,
             patch("vla_inference_server.snapshot_download") as download,
             patch(
                 "vla_inference_server.make_pre_post_processors",
@@ -712,7 +764,10 @@ def test_policy_runner_legacy_call_uses_existing_mutable_defaults(self) -> None:
             PolicyRunner("acme/model", "pi05", "cpu", 8, "EXP", 8)
 
         download.assert_not_called()
-        policy_class.from_pretrained.assert_called_once_with("acme/model")
+        read_config.assert_called_once_with("acme/model")
+        policy_class.from_pretrained.assert_called_once_with(
+            "acme/model", config=policy_config
+        )
         self.assertEqual(processors.call_args.kwargs["pretrained_path"], "acme/model")
 
     def test_pi05_loader_reads_weights_from_a_snapshot_directory(self) -> None:
@@ -820,6 +875,52 @@ def test_unknown_schedule_names_the_valid_values(self) -> None:
         self.assertIn("vla_serving.yaml", str(ctx.exception))
 
 
+class TestTrimVocabularyHeads(unittest.TestCase):
+    """trim_vocabulary_heads: drops the two heads and nothing else."""
+
+    @staticmethod
+    def build_model() -> torch.nn.Module:
+        """A stand-in with pi0.5's attribute path and a weight either side of it."""
+
+        def branch() -> torch.nn.Module:
+            part = torch.nn.Module()
+            part.lm_head = torch.nn.Linear(4, 8)
+            part.layers = torch.nn.Linear(4, 4)
+            return part
+
+        expert = torch.nn.Module()
+        expert.paligemma = branch()
+        expert.gemma_expert = branch()
+        model = torch.nn.Module()
+        model.paligemma_with_expert = expert
+        model.action_out_proj = torch.nn.Linear(4, 8)
+        return model
+
+    def test_both_heads_are_dropped(self) -> None:
+        model = self.build_model()
+        trim_vocabulary_heads(model)
+        self.assertIsNone(model.paligemma_with_expert.paligemma.lm_head)
+        self.assertIsNone(model.paligemma_with_expert.gemma_expert.lm_head)
+
+    def test_every_other_weight_survives_unchanged(self) -> None:
+        """The heads are the whole edit, so a chunk reads the same weights it did.
+
+        Actions leave through action_out_proj, which this reaches past; a trim
+        that touched anything on that path would change what the model commands.
+        """
+        model = self.build_model()
+        before = {
+            name: tensor.clone()
+            for name, tensor in model.state_dict().items()
+            if "lm_head" not in name
+        }
+        trim_vocabulary_heads(model)
+        after = model.state_dict()
+        self.assertEqual(sorted(after), sorted(before))
+        for name, tensor in before.items():
+            self.assertTrue(torch.equal(after[name], tensor), name)
+
+
 class TestDecodeImageB64(unittest.TestCase):
     """decode_image_b64: base64 JPEG -> CHW float32 [0,1] RGB tensor."""
 
@@ -937,6 +1038,7 @@ def __init__(
         native_map: dict | None = None,
     ) -> None:
         self.device = "cpu"
+        self.int8 = False
         self._infer_error = infer_error
         # Like PolicyRunner, derived once at construction.
         self.request_names = request_camera_names(