Add an Action Primitive#

An action primitive in RPent turns a tool call into an action that the environment can execute. It can be a learned policy (a VLA, a WAM, a diffusion planner) or a scripted routine (move_to, open_gripper). This page explains how to add either type.

Two Types of Primitives#

Family

Execution location

Examples

Model-based (VLA / WAM / diffusion / …)

Runs in its own process (vla_server) and is called through a model client held by the toolkit.

Pi0.5 (LIBERO), RLDX-1 (RoboCasa)

Scripted (kinematic / heuristic)

Runs in the agent process, with an optional server-side RPC for kinematics. It does not load model weights.

move_to, rotate_wrist, release, back_project

From the LLM’s perspective, both types expose the same interface: a tool schema, a primitives method, and a state dump after the call. They differ only in how the method is implemented.

Add a Scripted Primitive#

Declare the method on the robot’s primitives class and register its bound tool:

from rpent.tools import ToolResult, iter_tools, tool

class MyRobotPrimitives:
    @tool
    def open_drawer(self, dx: float = 0.15) -> ToolResult:
        """Pull the grasped drawer handle backwards.

        Args:
            dx: Pull distance in meters.
        """
        self._env.step(build_open_drawer_chunk(dx))
        return ToolResult(data={"ok": True, "dx": dx})

# In the toolkit, after constructing self._primitives:
self.add_tools(iter_tools(self._primitives))

Type annotations and the docstring define the model-facing schema. Toolkit validates arguments and captures state through get_env_state after execution. Use @tool(readonly=True) for readers such as view_env_state and back_project. For module-level tools, exclude internal resources with @tool(exclude=("state",)) and bind them with Tool.with_handler. All three planners use the registered declarations.

Add a VLA (or Other Model-based Primitive)#

Because the model runs in its own process, adding a model-based primitive requires a few additional components:

  1. Write ``vla_server.py``. This process owns only the model weights and CUDA context. Use rpent.robots.components.vla_facade_base.BaseVLAFacade as the base class, implement predict, and register any additional model RPCs by extending _register_rpc:

    • The default transport is HTTP (JSON over POST /call), which works well for flat image + state payloads such as the LIBERO / Pi0.5 pattern.

    • Switch to socket RPC (--transport socket) if your obs is a nested dict of numpy arrays with history stacks (avoids the JSON re-encode overhead).

    BaseVLAFacade registers vla.predict and serializes model calls; its inherited RpcFacade.serve handles transport binding, healthz, shutdown, parent-death detection, and resource cleanup.

  2. Write a model client. Subclass rpent.robots.components.vla_client_base.BaseVLAClient, which provides the common vla.predict call, and add only the environment-specific input / output adaptation. See rpent.robots.components.pi05_vla_client.Pi05VLAClient for the LIBERO implementation.

  3. Add a method to the primitives. In the current robot’s primitives class, call the model client, pass the returned action chunk to the environment, and return a ToolResult. The model client API is rpent.robots.components.pi05_vla_client.Pi05VLAClient.predict(), which reads the instruction from env_obs["task_descriptions"] and returns a [chunk, action_dim] numpy action chunk (batch dim already stripped):

    @tool
def mymodel_pick(self, target: str) -> ToolResult:
    """Pick the requested target using the model."""
    env_obs = self._env.get_obs()
    env_obs["task_descriptions"] = f"pick {target}"
    chunk = self._model.predict(env_obs)
    self._env.chunk_step(chunk)
    return ToolResult(data={"model": "mymodel", "target": target})

  4. Register the declared tool in the toolkit. Follow the same pattern as for a scripted primitive.

  5. Wire the components together in ``robot_spec.py``. The robot’s get_toolkit builds the toolkit with runtime_kwargs:

    def get_toolkit(*, runtime_kwargs, dashboard_events):
    from robots.myrobot.toolkit import MyRobotToolkit
    return MyRobotToolkit(
        runtime_kwargs=runtime_kwargs,
        dashboard_events=dashboard_events,
    )

    The robot package’s _init_runtime builds runtime_kwargs, for example {"env": MyRobotEnvClient(...), "model": MyModelClient(...)}. The toolkit constructor then forwards it to the primitives.

Reuse an Existing vla_server Across Runs#

Model servers often take a long time to start, so the runner can connect to an instance that is already running:

rpent --robot libero --vla-endpoint http://vla-host:8000 ...

If the model keeps per-episode state, expose a vla_reset RPC and call it between tasks. The same server process can then be reused safely across sequential runs.

Session-aware VLA Backends (per-client Policy State)#

Most VLA backends are stateless: predict only runs inference and keeps no per-client state, so session_id can be ignored. Some models do carry per-client policy state (e.g. RLDX-1’s memory/RTC); when a single vla_server serves multiple clients, their policy state would cross-contaminate, so it must be isolated per session. Wiring it up in three parts:

  • Facade side: construct the BaseVLAFacade subclass with enable_sessions=True and session_timeout_s, and implement _on_session_drop — clean up that client’s policy state when the session ends (the client’s session.close RPC or idle expiry). If you need an explicit reset, expose an extra reset_session RPC (clears policy state only, does not destroy the session). serve must pass session_sweep_s (> 0) so a background thread periodically reclaims expired sessions.

  • Client side: construct the RpcClient inside the model client with enable_sessions=True; it registers a session with the server on connect. session_id is derived from the connection and injected into the server-side handler by the facade — the client does not pass it, and must not forge session_ids inside predict’s options.

  • Primitives side: call reset_session before a task starts to clear policy state left over from the previous episode, so consecutive runs do not leak state into each other.

Single-threaded Serve (EGL-rendering Backends)#

Most backends use the serve inherited from their base class, which spawns a worker thread per request. If your server process renders with EGL (e.g. robosuite / MuJoCo offscreen rendering, see render_camera), the EGL context must stay on one thread, and concurrent dispatch would break context affinity.

Mix MainThreadServeMixin into your facade class (before BaseEnvFacade / BaseVLAFacade) and inherit the serve it overrides — it runs the transport server on a daemon thread but executes every dispatch serially on the thread that called serve (normally the process main thread), handing requests from the transport thread over via a work queue:

from rpent.utils.rpc.main_thread_serve import MainThreadServeMixin
from rpent.robots.components.env_facade_base import BaseEnvFacade

class MyEnvFacade(MainThreadServeMixin, BaseEnvFacade):
    ...

facade.serve(transport="http", host=host, port=port)  # dispatch on the main thread

The overridden serve keeps the same contract as RpcFacade’s serve: it still supports healthz / shutdown, parent-watch, and sessions (when constructed with enable_sessions=True, serve still requires session_sweep_s). Subclasses do not need to override serve to delegate — just inherit it (see RoboCasaEnvFacade in robots/robocasa/env_server.py). Backends that do not need EGL single-threading keep the plain inherited serve.

Design Principles for a New Primitive#

  • Tools describe intent, not motion. A good tool name is pi0_pick, not execute_action_chunk_of_length_20.

  • Stateful tools capture state before returning to the planner. Toolkit calls get_env_state after the handler while retaining exclusive execution. readonly tools skip this capture and may run alongside other readonly tools.

  • Keep result data small. Tool return values are fed back to the LLM as text. Save larger observations through EnvState.save; EnvState automatically records each logical base name in its owned StepRecord.artifacts set. Expose images through view_env_state and geometry through environment tools rather than returning raw paths.

  • Guardrails belong in env_server, not in the toolkit. The LLM can and will call any tool with any arguments; workspace bounds and safety clamps must be enforced on the server side.

Beyond VLAs#

The same pattern extends to non-VLA model primitives:

  • World Action Models (WAM) — imagination-based rollouts that produce a plan the env then executes. Wire them exactly like a VLA: their own process, their own client.

  • Diffusion planners / MPC — same shape; the “action” the tool returns may be a trajectory rather than a single chunk, and the env_server steps it out.

  • Multiple primitives sharing one server — a single vla_server can host several models; the tool decides which head to call via a model kwarg on predict.

Regardless of the implementation, the framework contract remains unchanged: model process → model client → primitives method → tool schema → Toolkit.add_tool.