Model Reference

Look up prediction methods, batching, action chunks, lifecycle hooks, and serving behavior here. Start with the Models guide if you are building a model; use the Python API for exact signatures.

The four corners

A model implements one or more of four predict corners. They sit on a 2×2 lattice over a batch axis (one lane vs. all vectorized lanes fused into one forward) and a chunk axis (one action vs. a chunk of future actions per forward).

On the chunk corners the leading axis of the return is the chunk axis. On the batch corners the leading axis is the batch axis; a chunk corner that is also batched returns the batch axis first, then the per-lane chunk axis ([N, H, ...]). A predict that subclasses Model and does not override one of these inherits a method that raises, so the runtime knows it is undefined.

Corner synthesis

Implement the most general corner your policy supports and the runtime fills the rest. It only ever derives downward along the lattice.

So predict_chunk_batch alone yields all four corners. Going up either axis is impossible: chunking is a model capability rather than glue, and batching up is left to the engine’s per-lane loop.

Two cases are worth pinning down. When both predict_batch and predict_chunk exist but predict does not, the runtime derives predict from the un-chunked predict_batch, so a single-step call stays un-chunked instead of paying for a chunk decode it would discard.

The native backend also cannot de-chunk: de-chunking indexes the chunk axis on array leaves, which the native rlmesh.Model over raw values does not have. A chunk-only model on that backend must define predict() explicitly, or construction raises a clear TypeError. The numpy, torch, and jax backends have no such restriction.

The predict context

Every corner may declare a trailing parameter named context — the name is what the runtime matches on, so an unrelated extra optional parameter is never mistaken for it. A corner that does not declare one is never handed one and pays nothing: no store entry, no allocation.

The per-lane corners receive one PredictContext; the batched corners receive a BatchPredictContext — a list with one entry per row, in batch order, because a fused batch may span independent episodes.

predict_index, predict_seed, and state come from the model’s own bounded episode store. Entries are dropped at the episode-end edge (the same edge that fires reset) and never otherwise: a live episode’s state is not evicted to make room. If ends never arrive, the store still cannot grow without bound: past 65,536 live episodes a new episode’s predict fails with a RuntimeError naming the ceiling. A deployment that legitimately runs more concurrent context-aware episodes than that raises it with RLMESH_EPISODE_CAPACITY.

rlmesh.predict_seed(episode_seed, predict_index) is the mixing law itself (FNV-1a, masked to 32 bits), exported so a sampler can build its own generator. Local and served drives compute it from the same function, so a seeded policy reproduces across both.

The execution horizon

execution_horizon is optional on the chunk corners and detected by arity, counting every positional parameter except a trailing context. A corner with a single positional parameter (predict_chunk(obs), or predict_chunk(obs, context)) is called without the horizon; a corner that declares a second one (predict_chunk(obs, execution_horizon=1)) receives it.

def predict_chunk(self, observation, execution_horizon=1):
    return self.policy.decode(observation, steps=execution_horizon)   # [execution_horizon, A]

Do not slice a fixed-size chunk down to the horizon yourself. Returning chunk[:execution_horizon] from a head that always produces the same length throws away actions the runtime was about to execute, and the runtime cannot tell that from a genuinely shorter chunk. Return the whole thing; the runtime takes the prefix.

Declaring your native chunk

Set native_chunk when your chunk length K is fixed – as a class attribute, or in load() when the checkpoint decides it:

class Chunker(rlmesh.torch.Model):
    native_chunk = 50

    def predict_chunk(self, observation):
        return self.policy.decode(observation)     # exactly 50 actions

class FromCheckpoint(rlmesh.torch.Model):
    def load(self):
        self.policy = load_policy()
        self.native_chunk = int(self.policy.horizon)

native_chunk is read after load() and reported when the adapter resolves. An execution_horizon above K is rejected during resolution. Each chunk prediction must return exactly K actions.

Leave native_chunk unset for variable-length output. The runtime then executes min(len(chunk), execution_horizon) actions and warns once if a returned chunk is shorter than the requested horizon.

Replay semantics

The runtime owns the replay. It calls the model once, splits the returned chunk along its leading axis, executes one action per environment step, and re-plans only when the queue drains. A horizon of 1 is a passthrough that never queues. A horizon above 1 caps the replay at min(len(chunk), execution_horizon) actions, so a receding-horizon model may emit a longer chunk than it re-plans. The horizon is bounded at 1024. The episode boundary (reset) drops any un-replayed tail, so a chunk never bleeds across episodes.

The chunk split treats a string, a bytes value, a mapping (a Dict-space action with the chunk axis inside each leaf), or a non-iterable scalar as a single-step chunk, matching the native split_chunk. An array’s leading axis is the chunk axis. A returned empty chunk raises rather than running an empty step.

execution_horizon only engages on a chunk corner. Requesting it on a model with no chunk corner warns and runs un-chunked, so the default of 1 is always safe. The end-to-end story (where the horizon is passed and how the rollout applies it) is in Running Evaluations.

Batched observation fusion

The batch corners do not receive a list of N observations. The runtime fuses the per-lane observations into one batched observation: every leaf gains a leading batch axis, so a Dict observation arrives as {key: array[N, ...]}, the shape every RL/VLA runtime hands a policy. You return the batched action the same way, and the runtime splits the batch axis back per lane.

def predict_batch(self, observations):
    # observations["image"] is array[N, H, W, C]; one forward over the batch.
    return self.net(observations["image"])          # return array[N, action_dim]

This is the tree_stack / tree_unstack machinery in _value_conversion.py. tree_stack recurses the container and stacks aligned leaves with the framework’s stack op; tree_unstack is its inverse and splits only the batch axis, leaving each lane’s chunk axis intact.

Device and framework handling

For torch and jax models, self.device is the single source of truth for placement. Set it inside load() when you move your weights:

def load(self):
    self.device = "cuda"
    self.policy = Policy.from_pretrained("org/checkpoint").to(self.device)

The backend itself only changes value conversion at the predict seam. The four corners and the lifecycle are identical across native, numpy, torch, and jax. See Framework Backends for choosing one and Framework Backends for the backend helpers.

Model quirks

Real checkpoints rarely speak the spec’s keys and shapes directly. The seams above absorb the difference so model code stays clean.

The _obs() remap helper is the SmolVLA pattern: the adapter delivers the payload under the spec’s keys, and a private method renames them to whatever dict the underlying policy expects.

def _obs(self, observations):
    return {
        "image": observations["observation.images.image"],
        "wrist": observations["observation.images.wrist"],
        "state": observations["observation.state"],
        "task": observations["instruction"],  # a list of N strings in a batch corner
    }

Tokenization stays in the model on purpose. Text delivers the instruction as a string; tokenize it inside your prediction function. There is no TokenizerInput on the spec.

Construction and lifecycle

The four lifecycle seams fire identically on the local loop and the served path.

There is no episode-begin hook. Per-episode state is lazy-seeded on the first predict, so a stateful model clears its state at episode end.

The served path drives the edge from the explicit ResetAdapter op, which names every id that ended, so on_episode_end fires once per id. Locally the session mints an id per reset() and fires the same hook at the next reset (and at close). Either way the model’s episode store drops that episode’s entry first.

Other construction inputs:

For a baseline that ignores observations, pass rlmesh.RANDOM_SAMPLE as the model to rlmesh.run / rlmesh.session. It samples the environment’s action space and resolves no adapter, even on a tagged env.

Serve and connect

serve(address) hosts the model as a blocking endpoint. A spec’d model resolves its adapter per env from the contract the runtime delivers in the handshake, then applies it around predict; a spec-less or NO_ADAPTER model serves its own predict directly.

MyPolicy().serve("0.0.0.0:50051")

Use the standard CLI to serve a model class from an importable module:

python -m rlmesh.serve my_pkg:MyPolicy

The serve env vars match the environment serve CLI:

The transport is plaintext gRPC and an endpoint is unauthenticated unless a bearer token is configured, so bind 0.0.0.0 only where the port is published to a trusted network. See Exposure and authentication.

Connect to a served model in one of two ways:

• rlmesh.RemoteModel(address): connect to a server you started.
• rlmesh.SandboxModel("image://my-model:latest"): start a local container from a prebuilt image. The image can be present locally or in an accessible registry. This does not launch a managed evaluation.

sess = rlmesh.session(rlmesh.SandboxModel("image://my-model:latest"), env)

Both handles plug into rlmesh.run / rlmesh.session like a local model. For a complete environment/model pair, see Bring Your Own Container; the sandbox mechanics are in Sandbox Environments, and connecting clients in Connect a Remote Environment.

Match your model’s shape

Find the row that matches your policy, then build it.

Common pitfalls