Trait PredictFn 

pub trait PredictFn: Send + Sync {
    // Required methods
    fn predict(
        &self,
        model_input: Value,
        episode: Option<&EpisodeInfo>,
    ) -> Result<Value>;
    fn predict_spec_less(
        &self,
        observation: ModelObservation,
    ) -> Result<Vec<SpaceValue>>;

    // Provided methods
    fn predict_chunk(
        &self,
        _model_input: Value,
        _execution_horizon: u32,
        _episode: Option<&EpisodeInfo>,
    ) -> Result<Option<Value>> { ... }
    fn has_chunk(&self) -> bool { ... }
    fn predict_batch(
        &self,
        _inputs: Vec<Value>,
        _episodes: &[EpisodeInfo],
    ) -> Result<Vec<Value>> { ... }
    fn has_batch(&self) -> bool { ... }
    fn predict_chunk_batch(
        &self,
        _inputs: Vec<Value>,
        _execution_horizon: u32,
        _episodes: &[EpisodeInfo],
    ) -> Result<Vec<Value>> { ... }
    fn has_chunk_batch(&self) -> bool { ... }
    fn predict_spec_less_chunked(
        &self,
        observation: ModelObservation,
        _execution_horizon: u32,
    ) -> Result<PredictFrames> { ... }
    fn allow_fusion(&self) -> bool { ... }
    fn native_chunk(&self) -> Option<u32> { ... }
    fn on_episode_end(&self, _episode_id: &str) -> Result<()> { ... }
    fn on_close(&self) -> Result<()> { ... }
}

The model’s predict callable plus its discovered lifecycle hooks.

predict is the contract floor: one already-assembled model input → one raw action. The engine loops it per lane (single-sample) for a spec’d route and runs apply_actions after. A spec-LESS route (no adapter) bypasses the engine via predict_spec_less, which gets the raw observation and preserves the pre-relocation batched path exactly.

Methods take &self: the model’s per-episode state lives in the host model object (e.g. a Python policy), not in this Rust handle, so a shared reference suffices and the engine can call back from a blocking worker thread.

Required Methods

fn predict( &self, model_input: Value, episode: Option<&EpisodeInfo>, ) -> Result<Value>

Spec’d route: one lane’s assembled model input → one raw action. The engine has already frame-stacked / customs’d / enc-shimmed the input. The input is a Value tree (a Map/List/leaf payload), matching the model spec’s InputNode shape — a bare tensor, a dict, or a tuple.

episode is this lane’s episode identity and explicit reset seed. On the fused grouped path (lanes from independent episodes batched into one forward) there is no single-episode identity: episode is None.

fn predict_spec_less( &self, observation: ModelObservation, ) -> Result<Vec<SpaceValue>>

Spec-less route (no adapter): the whole observation goes straight to the model, batched, returning one action per lane. Preserves the pre-engine behavior byte-for-byte (the binding reproduces the original path).

Provided Methods

fn predict_chunk( &self, _model_input: Value, _execution_horizon: u32, _episode: Option<&EpisodeInfo>, ) -> Result<Option<Value>>

Single-sample CHUNK corner: one assembled model input → a chunk of raw actions (the leading axis is the chunk axis, unstacked by split_chunk). None (the default) means the model has no distinct chunk corner, so the engine falls back to predict. A model that authors a separate chunk policy (e.g. a Python predict_chunk) returns Some(chunk).

execution_horizon is how many actions the runtime will execute before re-planning (pinned at ResolveAdapter). The model returns its native chunk; the engine executes a prefix of it (split_chunk(...).take(h)) and discards the rest, so a fixed-size head ignores the value and is correct either way. An autoregressive head may decode exactly execution_horizon actions to avoid wasting decode on a longer natural chunk.

episode follows the same contract as predict: real identity per lane, None on the fused grouped path.

fn has_chunk(&self) -> bool

Whether this model defines a chunk corner. Queried once at ResolveAdapter so the engine can warn when the runtime pins a horizon > 1 but the model cannot chunk (chunking is then inactive — the runtime re-plans every step).

fn predict_batch( &self, _inputs: Vec<Value>, _episodes: &[EpisodeInfo], ) -> Result<Vec<Value>>

Batched corner: N assembled lane inputs → N raw actions (one per lane) in a single call, so the model runs one forward pass for the whole vector. The engine prefers this over the per-lane predict loop when has_batch is true. Default unimplemented (only ever called when the flag is set).

episodes is row-aligned with inputs: row i’s identity and reset seed. A fused grouped batch concatenates lanes from independent episodes, so the list is per-row rather than per-call — there is no single episode identity here. Defaulted so an out-of-tree impl that only overrides predict keeps compiling.

fn has_batch(&self) -> bool

Whether this model defines the batched corner (predict_batch).

fn predict_chunk_batch( &self, _inputs: Vec<Value>, _execution_horizon: u32, _episodes: &[EpisodeInfo], ) -> Result<Vec<Value>>

Batched chunk corner: N assembled lane inputs → N native action chunks (leading axis = chunk) in a single call. Preferred for a vectorized chunked route when has_chunk_batch is true. execution_horizon is the runtime’s execution prefix (see predict_chunk); episodes is row-aligned with inputs (see predict_batch). Default unimplemented (gated by the flag).

fn has_chunk_batch(&self) -> bool

Whether this model defines the batched chunk corner (predict_chunk_batch).

fn predict_spec_less_chunked( &self, observation: ModelObservation, _execution_horizon: u32, ) -> Result<PredictFrames>

Spec-less route with a pinned execution horizon: like predict_spec_less, but when the runtime pinned execution_horizon > 1 a chunking binding runs its chunk corner on the raw observation and returns frame 0 plus the replay frames. The default ignores the horizon and stays un-chunked (a non-chunking model re-plans every step, exactly as before).

fn allow_fusion(&self) -> bool

Whether this model permits the fused forward pass: a grouped predict may concatenate lanes from different routes into ONE batched corner call. Default-OFF (a hand-written batched corner is not necessarily lane-independent); a binding whose batched corners fuse independent lanes by construction turns it on.

fn native_chunk(&self) -> Option<u32>

The model’s NATIVE chunk length K: how many per-step actions ONE chunk corner call returns, when the model declares it. None (the default) is the elastic contract — the model returns whatever it returns and the engine takes the min(len, execution_horizon) prefix.

Declaring K is a promise the engine holds the model to: the resolve doors reject execution_horizon > K (the runtime would replay frames the model never produced) and a chunk corner that returns anything other than exactly K frames fails the predict rather than silently short-replaying.

fn on_episode_end(&self, _episode_id: &str) -> Result<()>

Fires when an episode ends (structurally-discovered model hook), driven by the explicit ResetAdapter op — once per id it lists, so a model keyed by episode_id drops exactly the episode that ended. The engine separately evicts that episode’s frame buffers. There is no episode-begin hook: per-episode state is lazy-seeded on first predict, so a stateful model resets its state here at episode end rather than at a (no-longer-signalled) begin.

fn on_close(&self) -> Result<()>

Fires once at shutdown (structurally-discovered model hook, e.g. free GPU). The engine separately clears all frame buffers.

Dyn Compatibility

This trait is dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety".

Implementors