Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 9 additions & 15 deletions examples/inference/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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.

Expand All @@ -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.
1 change: 1 addition & 0 deletions examples/inference/franka.json
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
65 changes: 37 additions & 28 deletions examples/inference/franka.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand Down Expand Up @@ -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
Expand All @@ -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 = []
Expand Down Expand Up @@ -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):
Expand All @@ -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",
Expand Down Expand Up @@ -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)
Expand Down
2 changes: 1 addition & 1 deletion examples/inference/requirements.txt
Original file line number Diff line number Diff line change
@@ -1 +1 @@
vlagents @ git+https://github.com/RobotControlStack/vlagents.git@lerobot
vlagents==0.3.0
Loading