docs(robot): document [N, T, A] infer contract and BatchedAgent ownership

Spell out on Model.infer/ainfer that implementations must keep the leading
batch dim N (ainfer indexes [0], BatchedModel scatters rows along it) and add
a one-line assert in LeRobotModel.infer. Document that BatchedAgent mutates the
passed-in agent in place, leaving it permanently batched.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: lukass16

hud/agents/robot/batching.py

@@ -104,13 +104,19 @@ class BatchedAgent(Agent):
 
     Requires an in-process batchable model; :class:`~hud.agents.robot.model.RemoteModel`
     is not supported (the OpenPI server protocol has no batched-request shape).
+
+    Takes ownership of ``agent``: it swaps ``agent.model`` for a :class:`BatchedModel` wrapper
+    in place (so the wrapper is shared by every per-run clone). The passed-in instance is
+    therefore permanently batched — hand :class:`BatchedAgent` a dedicated agent and don't
+    also use that same instance for direct, unbatched :class:`RobotAgent` rollouts.
     """
 
     def __init__(self, agent: RobotAgent, *, batch_size: int, max_wait_s: float = 0.05) -> None:
         if agent.model is None:
             raise RuntimeError("BatchedAgent needs agent.model set")
         self._template = agent
-        # Wrap once; every per-run clone shares this batcher by reference.
+        # Wrap once, in place: the passed-in agent is now permanently batched (see class doc).
+        # Every per-run clone shares this batcher by reference.
         agent.model = BatchedModel(agent.model, batch_size=batch_size, max_wait_s=max_wait_s)
 
     async def __call__(self, run: Run, **kwargs: Any) -> None:

hud/agents/robot/model.py

@@ -30,11 +30,18 @@ class Model:
     """
 
     def infer(self, batch: Any) -> ActionArray:
-        """runs policy on a batch, returns [N, T, A] action chunk"""
+        """Run the policy on an ``[N, ...]`` batch, return an ``[N, T, A]`` chunk.
+
+        Implementations MUST keep the leading batch dim ``N`` (even for ``N == 1``):
+        :meth:`ainfer` indexes ``[0]`` and :class:`~hud.agents.robot.batching.BatchedModel`
+        scatters rows along it, so a squeezed ``[T, A]`` silently breaks both.
+        """
         raise NotImplementedError
 
     async def ainfer(self, batch: Any) -> ActionArray:
-        """Awaited single-rollout entry: run :meth:`infer` in a thread, return its ``[T, A]``."""
+        """Awaited single-rollout entry: run :meth:`infer` in a thread, return its single
+        ``[T, A]`` row. Indexing ``[0]`` assumes :meth:`infer` honors the ``[N, T, A]`` contract.
+        """
         return (await asyncio.to_thread(self.infer, batch))[0]
 
 
@@ -65,7 +72,9 @@ def infer(self, batch: Any) -> ActionArray:
         if self._first_inference:
             print("[agent] first inference done — inference is now fast", flush=True)
             self._first_inference = False
-        return chunk.float().cpu().numpy()
+        arr = chunk.float().cpu().numpy()
+        assert arr.ndim == 3, f"expected [N, T, A] chunk, got {arr.shape}"  # LeRobot keeps the N dim
+        return arr