Models

Note

This is the autodoc API reference. For the authoring guide see Models, and Model Reference for the full prediction-corner contract.

Model workers wrap a Python prediction function and run it against an RLMesh environment endpoint. The framework backend controls how observations are decoded before predict_fn runs and how returned actions are encoded.

Reach for a concrete Model class below in the value type your prediction function wants. Authors implement load() plus exactly one of the four prediction corners (predict, predict_chunk, predict_batch, or predict_chunk_batch); the runtime dispatches to whichever is defined.

Base Model

class rlmesh._models.base.ModelBase

Bases: Generic[ObsT, ActT]

A model: a policy you serve or drive against an env.

Construct one two ways:

• Wrap a predict callable – Model(lambda obs: ...).

• Subclass Model and override predict (and optionally load for weight loading, reset/close lifecycle hooks, and the spec class attribute), then instantiate it – class P(Model): ... then P().

run(env, seeds=...) drives the model against an env and returns a typed RunResult – it resolves the adapter from the env’s published tags and this model’s spec, so predict works in the model’s own input/output format with no per-env glue. serve() hosts the model as an endpoint for the runtime to dial.

Parameters:

• source – A predict callable. Omit it when subclassing (predict is the source then).

• spec – Optional rlmesh.adapters.ModelSpec; makes this an adapted model. Pass rlmesh.NO_ADAPTER to explicitly skip adapter resolution. Overrides the spec class attribute when both are set.

• on_close (on_episode_end /) – Optional lifecycle callbacks (they override a subclass’s reset/close methods, and a wrapped policy’s). There is no episode-begin hook: per-episode state is lazy-seeded on first predict, so a stateful model clears its state at episode end via on_episode_end (a subclass’s reset() is wired here). This fires identically on the local run(env)/session loop and the served wire path (driven there by the explicit ResetAdapter), so local and served behave the same.

• trust_entrypoints – Allow module:callable custom-input entrypoints in a spec to be imported during adapter resolution.

Examples

>>> from rlmesh.numpy import Model
>>> result = Model(lambda observation: 0).run("127.0.0.1:5555", seeds=[0])
>>> result.mean_reward
0.0

params: ClassVar[ParamSpec | None] = None

Optional declared construction-parameter surface for load – the model counterpart of rlmesh.EnvFactory.params, presented/swept the same way (see rlmesh.params). Advisory today: a dashboard reads it via rlmesh.describe; binding it into load is gated on the served-load seam.

device: object | None = None

Compute device for model inputs (torch/jax only). Set it in load() alongside moving your weights – one source of truth – and rlmesh moves every obs tensor leaf onto it before predict(), so you never call .to(device) yourself. None (the default) leaves obs as decoded; ignored by the numpy / raw-Value model (no device concept).

classmethod describe()

Return this model’s full metadata envelope (see rlmesh.describe()).

Return type:

dict[str, Any]

__init__(source=None, *, spec=None, on_episode_end=None, on_close=None, trust_entrypoints=False)

Parameters:

• source (Callable[[...], object] | object | None)

• spec (object | None)

• on_episode_end (Callable[[], None] | None)

• on_close (Callable[[], None] | None)

• trust_entrypoints (bool)

Return type:

None

spec: object | None = None

a ModelSpec, NO_ADAPTER, or None. Set it as a class attribute when subclassing; the spec= kwarg overrides it per instance.

Type:

The model’s content

load(**kwargs)

Load weights into self (from_pretrained etc.); heavy imports here.

Optional subclass hook; a no-op by default. Called once during __init__ (subclass mode only) before the native worker is built.

Parameters:

kwargs (Any)

Return type:

None

predict(observation)

Map one observation to an action (or an action chunk per spec).

Override when subclassing Model; the default raises. A model built by wrapping a predict callable uses that callable instead of this method.

Parameters:

observation (ObsT)

Return type:

ActT

predict_chunk(observation)

Optional: map one observation to a CHUNK of actions (leading axis = chunk).

Override this alongside predict() when the policy emits an action chunk in one forward pass (ACT, diffusion, flow, VLA action heads). Return your model’s native chunk; the runtime owns the replay – it executes the first execution_horizon actions one per step without re-calling the model, then re-plans. Defining this method is the chunk capability: a single-action model leaves it unimplemented and the runtime re-plans every step.

Most policies ignore the execution horizon – their chunk length is fixed by the trained weights, and the runtime simply uses a prefix of the native chunk. An autoregressive decoder that can stop early may add an optional second parameter execution_horizon: int = 1 (keep the default, so it stays a compatible override of this one-arg base); the runtime fills it with how many actions it will execute, so the head can decode exactly that many instead of its full natural length. The default raises.

Parameters:

observation (ObsT)

Return type:

ActT

predict_batch(observations)

Optional: one batched forward over all N vectorized lanes.

Override (alongside predict()) to run a single forward pass for a vectorized route instead of one call per lane. The runtime fuses the N per-lane observations into one batched observation – every leaf gains a leading batch axis, so a Dict observation arrives as {key: array[N, ...]} (a NumPy/Torch/JAX array per leaf, stacked for this model’s framework), not a list of N dicts. Return the batched action the same way: one value whose leaves carry the leading batch axis (e.g. array[N, action_dim]); the runtime splits it back per lane. The engine prefers this corner for a vectorized route; the default raises and the route is driven per-lane through predict().

(The dependency-free rlmesh.Model over raw Value trees can’t fuse opaque tensors, so it instead receives the per-lane list and returns one.)

Parameters:

observations (ObsT)

Return type:

ActT

predict_chunk_batch(observations)

Optional: one batched forward returning an action CHUNK per lane.

The batched counterpart of predict_chunk(): receives the fused batched observation (leaves [N, ...]; see predict_batch()) and returns the batched native chunk – leaves [N, chunk, ...] (batch axis first, then the per-lane chunk axis). The runtime splits the batch axis back per lane, executes a prefix (execution_horizon) of each, then re-plans. The engine prefers it for a vectorized chunked route. Like predict_chunk(), an autoregressive decoder may add an optional execution_horizon: int = 1 second parameter to decode exactly that many. The default raises.

Parameters:

observations (ObsT)

Return type:

ActT

reset()

Optional: called at each episode boundary (no-op by default).

Return type:

None

close()

Optional: release resources at the end of a run (no-op by default).

Return type:

None

run(env_or_address, *, seeds=None, max_episodes=None, instruction=None, close_env=False, token='', execution_horizon=1, view=None)

Drive this model against an env and return a RunResult.

Resolves the adapter from the env’s tags and this model’s spec, then runs a per-episode loop. seeds gives a per-episode seed and sets the episode count unless max_episodes is given; instruction, when given, overrides every text input the spec declares on each step – at its placement in the input tree (bare-root, top-level, or nested) and in that input’s declared shape (a bare str, or [instruction] for a container='list' text input). env_or_address is an env object exposing reset/step (e.g. a RemoteEnv), an EnvFactory (built and tag-stamped, then driven locally), an object with an address, or a bare address string the loop dials.

execution_horizon (> 1) executes that many actions of each predicted chunk one per env step, re-planning every execution_horizon steps – only when this model defines predict_chunk(); otherwise it runs un-chunked.

Parameters:

• env_or_address (object)

• seeds (Sequence[int] | None)

• max_episodes (int | None)

• instruction (str | None)

• close_env (bool)

• token (str)

• execution_horizon (int)

• view (object)

Return type:

RunResult

session(env_or_address, *, instruction=None, close_env=False, token='', trust_entrypoints=None, execution_horizon=1, view=None)

Bind this model to an env and return a Session to drive by hand.

The manual counterpart of run(): drive reset / predict / step yourself, or call Session.run() to pump whole episodes. env_or_address is an env object, an EnvFactory, a remote-env handle, or an address string (see run()). execution_horizon (> 1) executes that many actions per predicted chunk, one per env step, when this model defines predict_chunk() (see run()).

Parameters:

• env_or_address (object)

• instruction (str | None)

• close_env (bool)

• token (str)

• trust_entrypoints (bool | None)

• execution_horizon (int)

• view (object)

Return type:

Session[ObsT, ActT]

serve(address, *, token='', options=None)

Host this model as an endpoint (blocking).

A spec’d model resolves its adapter per env from the env contract the resolve_adapter handshake delivers, then applies it around predict; a spec-less / NO_ADAPTER model serves its own predict directly.

Parameters:

• address (str)

• token (str)

• options (ServeOptions | None)

Return type:

None

run_local(env_address, *, token='')

Native worker loop against a remote env, until the env ends.

Runs the session to completion for its side effects. Telemetry is surfaced on the serving runtime via its on_telemetry hook, not returned here.

Parameters:

• env_address (str)

• token (str)

Return type:

None

run_local_for_episodes(env_address, *, token='', max_episodes)

Native worker loop against a remote env for a fixed episode count.

Runs for the requested episode count for its side effects; see run_local() for where telemetry is surfaced.

Parameters:

• env_address (str)

• token (str)

• max_episodes (int)

Return type:

None

Concrete Models

Concrete backend model classes inherit ModelBase and only change value conversion:

Class	Import	Observation type	Action encoding
Native model	rlmesh.Model	RLMesh-native values and primitives	RLMesh-native values
NumPy model	rlmesh.numpy.Model	NumPy arrays, primitives, and containers	NumPy arrays and primitives
Torch model	rlmesh.torch.Model	Torch tensors, primitives, and containers	Torch tensors and primitives
JAX model	rlmesh.jax.Model	JAX arrays, primitives, and containers	JAX arrays and primitives

See Framework Backends for backend helpers.