Models

Note

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

Models 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 predict corners (predict, predict_chunk, predict_batch, or predict_chunk_batch); the runtime dispatches to whichever is defined.

Base Model

class rlmesh._models.base.ModelBase[source]

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.

Example:

from rlmesh.numpy import Model

result = Model(lambda observation: 0).run("127.0.0.1:5555", seeds=[0])
print(result.mean_reward)

Build the model from source (see the class docstring for the modes).

The overloads are the typed face of the permissive runtime contract: a predict callable flows its annotated observation/action types onto the model’s generic parameters (Model(predict) with def predict(obs: X) -> Y is a Model[X, Y]), while a class source (instantiated by coerce_model()), a duck-typed policy object, or an unannotated callable falls back to the framework’s default value types.

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()[source]

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

Return type:

dict[str, Any]

__init__(source: type[Any], *, spec: object | None = None, on_episode_end: Callable[[], None] | None = None, on_close: Callable[[], None] | None = None, trust_entrypoints: bool = False) → None[source]
__init__(source: Callable[[ObsT], ActT], *, spec: object | None = None, on_episode_end: Callable[[], None] | None = None, on_close: Callable[[], None] | None = None, trust_entrypoints: bool = False) → None
__init__(source: object | None = None, *, spec: object | None = None, on_episode_end: Callable[[], None] | None = None, on_close: Callable[[], None] | None = None, trust_entrypoints: bool = False) → None

Build the model from source (see the class docstring for the modes).

The overloads are the typed face of the permissive runtime contract: a predict callable flows its annotated observation/action types onto the model’s generic parameters (Model(predict) with def predict(obs: X) -> Y is a Model[X, Y]), while a class source (instantiated by coerce_model()), a duck-typed policy object, or an unannotated callable falls back to the framework’s default value types.

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)[source]

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)[source]

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)[source]

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 execution_horizon: int = 1 parameter – as a second positional parameter or keyword-only (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)[source]

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)[source]

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()[source]

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

Return type:

None

close()[source]

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, max_episode_steps=None, max_episode_seconds=None, hooks=None, instruction=None, close_env=False, trust_entrypoints=None, execution_horizon=1, view=None)[source]

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). (A served model rejects instruction= – the env’s own instruction is used.) 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.

max_episode_steps / max_episode_seconds cap each episode (hitting a cap marks it truncated), and hooks (rlmesh.RunHooks) observes the loop; all three pass through to Session.run.

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.

This one-shot convenience creates its session internally and closes it (honoring close_env) when the run ends. To reuse one binding across several runs, hold a session() yourself – Session.run() leaves a caller-held session open.

Parameters:

  • env_or_address (LocalEnvTarget)

  • seeds (Sequence[int] | None)

  • max_episodes (int | None)

  • max_episode_steps (int | None)

  • max_episode_seconds (float | None)

  • hooks (RunHooks | None)

  • instruction (str | None)

  • close_env (bool)

  • trust_entrypoints (bool | None)

  • execution_horizon (int)

  • view (ViewArg)

Return type:

RunResult

session(env_or_address, *, instruction=None, close_env=False, trust_entrypoints=None, execution_horizon=1, view=None)[source]

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 – as many times as you like; the caller-held session (its connection, viewer, and adapter state) stays open until you close it (close() or the with block). 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 (LocalEnvTarget)

  • instruction (str | None)

  • close_env (bool)

  • trust_entrypoints (bool | None)

  • execution_horizon (int)

  • view (ViewArg)

Return type:

Session[ObsT, ActT]

serve(address, *, options=None)[source]

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:
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.

Served and Sandboxed Models

A model does not have to run in your process. RemoteModel dials a policy that is already served on an endpoint; SandboxModel runs a prebuilt image:// tag in its own container. Both bind to an environment through session() and expose the same Session drive loop.

final class rlmesh.RemoteModel[source]

Bases: RemoteModelBase[None | bool | int | float | str | bytes | _Tensor | list[Value] | tuple[Value, …] | dict[str, Value], None | bool | int | float | str | bytes | _Tensor | list[Value] | tuple[Value, …] | dict[str, Value]]

class rlmesh.SandboxModel[source]

Bases: object

A model served from an isolated container.

The source is a prebuilt rlmesh-serving image you built yourself (BYO), given as a bare smolvla:latest or an explicit image:///docker:// tag: the image is run directly and its baked CMD drives it. **params are the model’s construction params – forwarded into the container as the load(**binding) binding (RLMESH_MAKE_KWARGS), validated against the model’s declared params before weights load.

A SandboxRuntime (runtime=) sets docker run flags – gpus for CUDA compute, or devices=["nvidia.com/gpu=all"] (a CDI ref, full graphics+compute) and volumes=[...] for large local checkpoints/assets. A model is always a prebuilt image, so there is no build config (no SandboxBuild).

__init__(source, /, *, runtime=None, **params)[source]
Parameters:
Return type:

None

serve()[source]

Start a long-lived container serving the policy as a model endpoint.

The prebuilt image:// image is run in serve mode.

Idempotent: a second call returns the already-running handle. The endpoint is reachable at address until shutdown().

Return type:

SandboxModel

session(env, *, instruction=None, close_env=False, trust_entrypoints=None, execution_horizon=1, connect_timeout_seconds=30.0, view=None)[source]

Serve this model and bind it to env, returning a neutral rlmesh.Session.

The managed sibling of rlmesh.RemoteModel.session(): starts the model container (idempotent), then opens a route configured from the env’s contract so the same drive loop works for both pairs:

with rlmesh.SandboxEnv("gym://CartPole-v1") as env:
    sess = rlmesh.session(
        rlmesh.SandboxModel("image://my-model:latest"), env
    )
    obs, _ = sess.reset()
    while not sess.done:
        obs, reward, terminated, truncated, _ = sess.step(sess.predict(obs))

Closing the session stops the container it started. instruction / trust_entrypoints cannot be honored by a served container and raise ValueError when set rather than being silently dropped. view opts into the built-in live viewer over this session’s env loop (see rlmesh.run()). execution_horizon requests open-loop action chunking: the runtime executes that many actions of each predicted chunk before re-planning (1 = re-plan every step; only engages if the served policy defines a chunk corner). Retries the connection while the container starts, up to connect_timeout_seconds.

Parameters:

  • env (EnvTarget)

  • instruction (str | None)

  • close_env (bool)

  • trust_entrypoints (bool | None)

  • execution_horizon (int)

  • connect_timeout_seconds (float)

  • view (ViewArg)

Return type:

Session[Any, Any]

shutdown()[source]

Stop the served container, if any. Idempotent; safe to call repeatedly.

Return type:

None

Sandbox Configuration

SandboxBuild groups the build-from-source settings for a gym:// / hf:// env source; SandboxRuntime groups the docker run flags applied when a prebuilt container starts. See Sandbox Environments for where each option goes.

class rlmesh.SandboxBuild[source]

Bases: object

Build-from-source infrastructure for a sandbox image.

Configures how the image is built from a gym:///hf:// source (base image, extra packages, the rlmesh pin, …). Meaningless for a prebuilt image (run, not built) – those fields are ignored with a warning. Pair with SandboxRuntime for docker run settings. Env/model construction params stay in the sandbox’s **params (the make/load binding).

__init__(base_image=None, rlmesh_package=None, packages=None, imports=None, trust_remote_code=False, allow_unpinned_hf=False, build_memory=None)
Parameters:

  • base_image (str | None)

  • rlmesh_package (str | PathLike[str] | None)

  • packages (Sequence[str] | None)

  • imports (Sequence[str] | None)

  • trust_remote_code (bool)

  • allow_unpinned_hf (bool)

  • build_memory (str | None)

Return type:

None

class rlmesh.SandboxRuntime[source]

Bases: object

Container run-time settings for a sandbox – docker run flags.

Applies whenever the container is run (prebuilt images and build-then-run), unlike SandboxBuild which only configures the from-source build. The host-resource knobs sim envs/models need:

  • gpus: docker run --gpus ("all" / a count / "device=0,1") – legacy GPU injection (CUDA compute only). NOTE: this does not inject the graphics/Vulkan driver; for SAPIEN/Vulkan use devices with a CDI ref.

  • devices: docker run --device entries, e.g. ["nvidia.com/gpu=all"] (a CDI device ref – the full graphics+compute driver) or a /dev node.

  • volumes: docker run -v mounts, e.g. ["/host/assets:/ctr/assets", "/host/ro:/ctr/ro:ro"].

  • user: docker run --user passthrough – an int uid or a "uid[:gid]" string. Set it to the owner of writable volume mounts (typically os.getuid()); the container runs --cap-drop ALL, so even root inside it obeys the mount’s permission bits.

Build-from-source attaches none of these (it has no docker run step); set them only for a prebuilt image source.

__init__(gpus=None, devices=None, volumes=None, user=None)
Parameters:

  • gpus (str | int | None)

  • devices (Sequence[str] | None)

  • volumes (Sequence[str] | None)

  • user (str | int | None)

Return type:

None