Merge pull request #325 from RobotControlStack/juelg/bump-agents

bump(examples): vlagents interface refactor

Author: juelg

examples/inference/README.md

@@ -12,13 +12,12 @@ Before starting `franka.py`, make sure a `vlagents` policy server is already run
 
 The policy server setup and supported launch commands are documented in:
 
-- [RobotControlStack/vlagents](https://github.com/RobotControlStack/vlagents)
-- [vlagents/README.md](../../vlagents/README.md)
+- [vlagents](https://github.com/RobotControlStack/vlagents)
 
 Typical server startup looks like:
 
 ```shell
-python -m vlagents start-server lerobot --port 20000 --host 0.0.0.0 --kwargs '{"policy_name": "act", "checkpoint_path": "<path to pretrained_model>", "n_action_steps": 1}'
+uv run python -m vlagents start-server lerobot --port 20000 --host 0.0.0.0 --kwargs '{"policy_name": "act", "checkpoint_path": "<path to pretrained_model>"}'
 ```
 
 For other policies such as `pi05` or `xvla`, use the matching startup command from the `vlagents` README and make sure the values in `franka.json` point at that server.
@@ -31,12 +30,13 @@ For other policies such as `pi05` or `xvla`, use the matching startup command fr
 - `vlagents_port`: Port exposed by the policy server.
 - `vlagents_model`: Agent id passed to `vlagents`, for example `lerobot`.
 - `instruction`: Natural-language task instruction sent to the policy on reset.
-- `robot_keys`: Robot ordering used to pack observations and unpack actions. The script assumes one 8-value action block per robot in this order: `7` joint values plus `1` gripper value.
+- `robot_keys`: Robot names expected in each returned action dictionary and used to construct per-robot observations.
 - `jpeg_encoding`: Whether observations are sent to the policy server using JPEG-compressed images.
 - `on_same_machine`: Set this according to whether the policy server runs on the same machine as the control process.
+- `image_size`: Client-side `(width, height)` resize applied before JPEG or shared-memory transport; defaults to `[224, 224]`. Set it to `null` to retain native resolution.
 - `fps`: Control loop target frequency used by the local rate limiter.
 - `record_path`: Output directory used when recording episodes.
-- `n_action_steps`: If `null`, the script requests one action per control step. If set to an integer greater than `0`, the script buffers that many actions from each policy response chunk.
+- `n_action_steps`: Local action-chunk execution horizon. If `null`, the script requests and executes one action per control step. If set to a positive integer, it buffers up to that many actions from each policy response chunk.
 - `max_rel_mov_joints`: Maximum allowed relative joint movement per step when running in joint control mode.
 - `max_rel_mov_cart`: Maximum allowed relative Cartesian translation and rotation per step when running in Cartesian modes.
 
@@ -57,23 +57,17 @@ When [franka.py](franka.py) is running, it waits for keyboard input on stdin. Th
 
 The script translates RCS observations to the `vlagents` `Obs` format as follows:
 
-- Every camera frame in `obs["frames"]` is converted to RGB and resized to `224x224`.
-- State is built by iterating through `robot_keys` in order and concatenating each robot's `joints` and `gripper` values.
+- Camera frames are passed to `RemoteAgent` at native resolution; the client resizes them to `image_size` before JPEG or shared-memory transport.
+- Each robot gets a `SingleObs` containing the shared camera set plus its own joints and gripper state.
 
-Action decoding is also order-dependent:
-
-- For each robot in `robot_keys`, the script reads `8` values from the policy action vector.
-- Values `0:7` become the robot joint command.
-- Value `7:8` becomes the robot gripper command.
-
-That means `robot_keys` must match the policy's expected robot ordering exactly.
+Action chunks contain one action dictionary per environment step. For each robot, the script forwards `SingleAct.action` as the joint command and `SingleAct.gripper` as the gripper command. The action dictionary must include every configured `robot_key`.
 
 ## Running
 
 After the policy server is up and `franka.json` is configured, run:
 
 ```shell
-python examples/inference/franka.py
+uv run python examples/inference/franka.py
 ```
 
 If the policy server is unreachable, the script will keep retrying connection until it becomes available or you exit.

examples/inference/franka.json

@@ -9,6 +9,7 @@
   ],
   "jpeg_encoding": true,
   "on_same_machine": false,
+  "image_size": [224, 224],
   "fps": 30,
   "record_path": "inference_recordings_bin_sort_duobench_xvla_bin_sort_real_2026-05-20_23-25-47_040000",
   "n_action_steps": 30,

examples/inference/franka.py

@@ -10,7 +10,6 @@
 
 import gymnasium as gym
 import numpy as np
-from PIL import Image
 from rcs._core.common import BaseCameraConfig, RobotPlatform
 from rcs._core.sim import SimConfig
 from rcs.envs.base import ControlMode, RelativeTo
@@ -20,7 +19,7 @@
 
 # from rcs_duobench.tasks.bin_sort import BinSortEnvConfig
 from vlagents.client import RemoteAgent
-from vlagents.policies import Act, Obs
+from vlagents.policies.interface import Obs, SingleAct, SingleObs
 
 import rcs
 
@@ -103,6 +102,7 @@ class InferenceConfig:
     robot_keys: list[str] = field(default_factory=lambda: ["left", "right"])
     jpeg_encoding: bool = True
     on_same_machine: bool = False
+    image_size: tuple[int, int] | None = (224, 224)
     fps: int = FPS
     record_path: str = RECORD_PATH
     n_action_steps: int | None = None
@@ -126,7 +126,12 @@ def __init__(self, env: gym.Env, cfg: InferenceConfig):
         self._command_queue: Queue[str] = Queue()
         self._shutdown_requested = threading.Event()
         self.remote_agent = RemoteAgent(
-            cfg.vlagents_host, cfg.vlagents_port, cfg.vlagents_model, cfg.on_same_machine, cfg.jpeg_encoding
+            cfg.vlagents_host,
+            cfg.vlagents_port,
+            cfg.vlagents_model,
+            cfg.on_same_machine,
+            cfg.jpeg_encoding,
+            cfg.image_size,
         )
         self.frame_rate = SimpleFrameRate(self._cfg.fps)
         self._action_buffer = []
@@ -164,41 +169,45 @@ def _drain_commands(self) -> tuple[bool, bool, bool, bool, bool]:
         return start_requested, record_requested, success_requested, stop_requested, reload_requested
 
     def obs_rcs2agents(self, obs: dict, info: dict | None = None) -> Obs:
-        cameras = {}
-        for frame in obs["frames"]:
-            cameras[frame] = obs["frames"][frame]["rgb"]["data"]
-            cameras[frame] = np.array(Image.fromarray(cameras[frame]).resize((224, 224), Image.Resampling.BILINEAR))
+        cameras = {frame: obs["frames"][frame]["rgb"]["data"] for frame in obs["frames"]}
 
-        state = []
+        obs_by_robot = {}
         for robot in self._cfg.robot_keys:
-            # TODO: currently hardcoded for joints
-            state.append(obs[robot]["joints"])
-            state.append(obs[robot]["gripper"])
+            obs_by_robot[robot] = SingleObs(
+                cameras=copy.deepcopy(cameras),
+                joints=np.asarray(obs[robot]["joints"], dtype=np.float32),
+                gripper=float(obs[robot]["gripper"]),
+                xyzrpy=np.asarray(obs[robot]["xyzrpy"], dtype=np.float32) if "xyzrpy" in obs[robot] else None,
+                tquat=np.asarray(obs[robot]["tquat"], dtype=np.float32) if "tquat" in obs[robot] else None,
+                # info=copy.deepcopy(info) if info is not None else {},
+            )
 
-        return Obs(cameras=cameras, gripper=None, info=info, state=np.concatenate(state))
+        return Obs(obs=obs_by_robot, language_instruction=self._cfg.instruction)
 
-    def act(self, obs_dict) -> None:
-        done = False
+    def act(self, obs_dict: Obs) -> dict[str, SingleAct]:
         if self._cfg.n_action_steps is None:
-            return self.remote_agent.act(obs_dict)
+            action_chunk = self.remote_agent.act(obs_dict).acts
+            if not action_chunk:
+                message = "Received empty action chunk from policy"
+                raise ValueError(message)
+            return action_chunk[0]
         if len(self._action_buffer) == 0:
             action = self.remote_agent.act(obs_dict)
-            selected_action = action.action[: self._cfg.n_action_steps]
-            self._action_buffer = selected_action.tolist()
-            done = action.done
+            selected_action = action.acts[: self._cfg.n_action_steps]
+            self._action_buffer = list(selected_action)
             if RELATIVETO == RelativeTo.CONFIGURED_ORIGIN:
                 for robot in self.env.get_wrapper_attr("envs"):
                     self.env.get_wrapper_attr("envs")[robot].get_wrapper_attr("set_origin_to_current")()
-        act = self._action_buffer.pop(0)
-        return Act(action=act, done=done)
+        return self._action_buffer.pop(0)
 
-    def action_agents2rcs(self, action: Act) -> dict[str, Any]:
+    def action_agents2rcs(self, action: dict[str, SingleAct]) -> dict[str, Any]:
         act = {}
-        for idx, robot in enumerate(self._cfg.robot_keys):
-            # TODO: this is currently hard coded for franka joints
-            act[robot] = {}
-            act[robot]["joints"] = action.action[idx * 8 : idx * 8 + 7]
-            act[robot]["gripper"] = action.action[idx * 8 + 7 : idx * 8 + 8]
+        for robot in self._cfg.robot_keys:
+            robot_action = action[robot]
+            act[robot] = {
+                "joints": np.asarray(robot_action.action, dtype=np.float32),
+                "gripper": np.asarray([robot_action.gripper], dtype=np.float32),
+            }
         return act
 
     def loop(self):
@@ -222,6 +231,7 @@ def loop(self):
                         model=self._cfg.vlagents_model,
                         on_same_machine=self._cfg.on_same_machine,
                         jpeg_encoding=self._cfg.jpeg_encoding,
+                        image_size=self._cfg.image_size,
                     )
                     logger.info(
                         "reloaded config from %s with host=%s port=%s model=%s",
@@ -269,14 +279,13 @@ def loop(self):
                         if record_requested:
                             self.env.start_record()
                     logger.info("starting episode%s", " with recording" if record_requested else "")
-                    self.remote_agent.reset(copy.deepcopy(obs_dict), instruction=self._cfg.instruction)
                     self._episode_running = True
                 else:
                     sleep(0.05)
                     continue
 
             action = self.act(copy.deepcopy(obs_dict))
-            if action.done:
+            if any(robot_action.done for robot_action in action.values()):
                 logger.info("done issued by agent, resetting environment")
                 obs, _ = self.env.reset()
                 obs_dict = self.obs_rcs2agents(obs)

examples/inference/requirements.txt

@@ -1 +1 @@
-vlagents @ git+https://github.com/RobotControlStack/vlagents.git@lerobot
+vlagents==0.3.0