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.

Declare tools from Python signatures#

@tool generates a tool’s description, JSON Schema, and argument validator from its signature and Google-style docstring. It works with plain functions and instance methods, so an existing primitives object can keep its clients and state:

from typing import Annotated

from pydantic import Field

from rpent.tools import ToolResult, iter_tools, tool

class MyPrimitives:
    def __init__(self, env):
        self.env = env

    @tool
    def move_delta(
        self,
        delta_xyz: Annotated[list[float], Field(min_length=3, max_length=3)],
    ) -> ToolResult:
        """Move the TCP by a base-frame offset.

        Args:
            delta_xyz: XYZ displacement in metres.
        """
        return ToolResult(data=self.env.move_delta(delta_xyz))

# In your robot toolkit, after super().__init__(...):
self._primitives = MyPrimitives(env)
self.add_tools(iter_tools(self._primitives))

Register the bound method from the instance. self is excluded from the schema, and each instance retains its own resources. Methods remain callable from Python, including through self.move_delta(...) and inherited methods. An existing undecorated method can also be registered with self.add_tool(tool(self._primitives.move_delta)).

Public parameters need type annotations and must accept keyword arguments. Annotated[..., Field(...)] supplies constraints. Signature defaults apply at runtime; advertise them in the schema only when intended, with Field(json_schema_extra={"default": value}). Toolkit.execute_tool applies Pydantic validation before executing the handler or capturing observations, rejecting unknown arguments and non-finite numbers. Direct Python calls retain normal Python argument handling.

Handlers return ToolResult(data=..., images=..., error=...). Put structured values in data, PNG bytes in images, and failures in error. A primitive that calls another tool directly receives the same ToolResult. Use @tool(readonly=True) to skip automatic observation capture. Tools declared with @tool or @tool() capture observations by default when called through Toolkit.execute_tool. Direct Python calls do not capture observations. readonly controls this capture step; it does not prohibit file writes or allow concurrent tool execution. iter_tools collects decorated members from the supplied instances or modules, including inherited methods. Adding a primitive requires decorating its method; there is no separate schema or tool-name list to update. Undecorated methods and properties are not collected. Toolkit still owns mode-specific filtering, resource binding, and execution guards. For an injected parameter such as state, declare exclude=("state",) and register declaration.with_handler(partial(declaration, state=self.state)). The decorator does not create environment or model clients.

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 ToolResult(data=...) with the action log. 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):

    def mymodel_pick(self, target: str) -> ToolResult:
    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. Decorate the method with ``@tool`` and register its bound method. Use type annotations and an Args docstring as in the example above.

  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.

  • Every tool ends with a state dump. The next turn depends on the state dump reflecting the post-action world. Don’t let the primitive return before the render finishes.

  • Keep ``ToolResult.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.