Unblock the DROID policy server setup on CUDA 13 (#187)

### Summary

Moves the **Cosmos3-Policy-DROID** server setup onto torch 2.13, adds
the instrumentation needed to measure its inference latency during
RoboLab evaluation, and makes the apex build optional — it currently
cannot build against the CUDA 13 base image at all, and is the slowest
layer by a wide margin.

### Changes

- **`docs/action_policy_droid_server.md`** — Sync the container with
`--group=cu130-torch213-train` rather than `cu130-train`, putting the
policy server on torch 2.13. The image itself still bakes in `cu130`
(torch 2.10), so the container re-syncs over the baked-in wheels at
startup. The documented `docker build` also passes `--build-arg
INSTALL_APEX=0`, since the policy server never imports apex. Also drops
the paragraph directing CUDA 12.x users to `cu128-train`.
- **`scripts/action_policy_server_robolab.py`** — Time each
`RobolabPolicyService.infer` call and print the elapsed milliseconds
when `EVAL_VERBOSE` is set, to measure server latency during RoboLab
evaluation runs. Gated behind the env var so evaluation runs opt in
without altering default server output.
- **`Dockerfile`** — Bump the apex pin `bf903a2` → `9e3568a` and add an
`INSTALL_APEX` build arg. Pin `bf903a2` assembles its own `-gencode`
flags and appends `arch=compute_70,code=sm_70` unconditionally, gated
only by lower bounds on the CUDA version, so several extensions fail on
the CUDA 13 base image with `nvcc fatal : Unsupported gpu architecture
'compute_70'`. Those code paths never consult `TORCH_CUDA_ARCH_LIST`, so
no arch override can work around it; upstream has since removed the
hand-built lists and derives the architectures from the detected CUDA
version instead. `INSTALL_APEX=0` skips the build entirely, which is
worth having regardless: apex is optional, its only import being the
`try`/`except`-guarded one in `callbacks/norm_monitor.py`.

---------

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Co-authored-by: lfengad <liangf@nvidia.com>

Author: 3 people

Dockerfile

@@ -52,11 +52,19 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     uv sync --locked --no-install-project --no-editable --all-extras --group=$(cat /root/.cuda-name) --group=vllm
 ENV PATH="/workspace/.venv/bin:$PATH"
 
+# Set to 0 to skip the apex build, which is by far the slowest layer. apex is optional:
+# the only import is in cosmos_framework/callbacks/norm_monitor.py, behind a try/except.
+ARG INSTALL_APEX=1
+
 # install apex (compiled C++/CUDA extensions; needs torch already present), so this line should be after the uv sync command.
 RUN --mount=type=cache,target=/root/.cache/uv \
-    VIRTUAL_ENV=/workspace/.venv APEX_CPP_EXT=1 APEX_CUDA_EXT=1 \
-    uv pip install -v --no-build-isolation --no-deps \
-        git+https://github.com/NVIDIA/apex@bf903a2
+    if [ "$INSTALL_APEX" = "1" ]; then \
+        VIRTUAL_ENV=/workspace/.venv APEX_CPP_EXT=1 APEX_CUDA_EXT=1 \
+        uv pip install -v --no-build-isolation --no-deps \
+            git+https://github.com/NVIDIA/apex@9e3568a; \
+    else \
+        echo "INSTALL_APEX=$INSTALL_APEX, skipping apex"; \
+    fi
 
 # Triton bundled ptxas doesn't support latest GPU architectures
 ENV TRITON_PTXAS_PATH="/usr/local/cuda/bin/ptxas"

cosmos_framework/scripts/action_policy_server_robolab.py

@@ -26,7 +26,9 @@
 init_script()
 
 import json
+import os
 import socket
+import time
 import threading
 from dataclasses import dataclass
 from pathlib import Path
@@ -574,6 +576,7 @@ def _build_sample(self, obs: dict[str, Any]) -> dict[str, Any]:
         return sample
 
     def infer(self, obs: dict[str, Any]) -> dict[str, Any]:
+        start_time = time.monotonic()
         sample = self._build_sample(obs)
         data_batch = _build_data_batch_from_sample(sample)
         seed = self._next_seed()
@@ -616,6 +619,11 @@ def infer(self, obs: dict[str, Any]) -> dict[str, Any]:
             video = self.model.decode(pred_vision_latent)  # [1,C,T,H,W]
             video = ((video[0].clamp(-1.0, 1.0) + 1.0) * 127.5).to(torch.uint8).permute(1, 2, 3, 0)  # [T,H,W,3]
             outputs["video"] = video.detach().cpu().numpy()
+
+        infer_time = time.monotonic() - start_time
+        infer_ms = infer_time * 1000.0
+        if os.environ.get("EVAL_VERBOSE"):
+            print(f"infer_ms: {infer_ms:.1f}")
         return outputs
 
 

docs/action_policy_droid_server.md

@@ -33,6 +33,7 @@ Build the Docker image:
 
 ```bash
 docker build \
+  --build-arg INSTALL_APEX=0 \
   -t cosmos-framework:latest \
   .
 ```
@@ -57,14 +58,12 @@ docker run \
   bash -c '\
     uv sync \
       --all-extras \
-      --group=cu130-train \
+      --group=cu130-torch213-train \
       --group=policy-server && \
     exec bash; \
   '
 ```
 
-The `--group=cu130-train` line targets CUDA 13.x drivers. On CUDA 12.x systems, replace it with `--group=cu128-train` (see the [Cosmos3 Cookbooks: Environment Setup](https://github.com/NVIDIA/cosmos/blob/main/cookbooks/cosmos3/README.md) for details).
-
 Inside the container, start the policy server:
 
 1. For [Cosmos3-Nano-Policy-DROID](https://huggingface.co/nvidia/Cosmos3-Nano-Policy-DROID), run: