rlmesh 0.1.0

//! The language-agnostic predict/resolve holes the served engine calls back
//! into.
//!
//! The vectorized stateful engine ([`AdaptedModelHandler`](super::engine::AdaptedModelHandler))
//! owns the per-lane loop, the episode-keyed frame buffers, and the native
//! adapter application — all pure Rust. The two genuinely host-language steps
//! cross via these traits, which a binding (PyO3, or any future language)
//! implements; a pure-Rust model implements them with no host runtime at all.

use super::types::EpisodeInfo;
use async_trait::async_trait;
use rlmesh_adapters::v1::{CustomTransform, EncodingTransform, ResolvedAdapter, Value};

use crate::spaces::{EnvContract, SpaceSpec, SpaceValue};
use crate::{Result, model::types::ModelObservation};

/// 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`](PredictFn::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.
pub trait PredictFn: Send + Sync {
    /// 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(&self, model_input: Value, episode: Option<&EpisodeInfo>) -> Result<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`](rlmesh_adapters::v1::split_chunk)). `None` (the default)
    /// means the model has no distinct chunk corner, so the engine falls back to
    /// [`predict`](Self::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`](Self::predict):
    /// real identity per lane, `None` on the fused grouped path.
    fn predict_chunk(
        &self,
        _model_input: Value,
        _execution_horizon: u32,
        _episode: Option<&EpisodeInfo>,
    ) -> Result<Option<Value>> {
        Ok(None)
    }

    /// 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 has_chunk(&self) -> bool {
        false
    }

    /// 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`](Self::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 predict_batch(&self, _inputs: Vec<Value>, _episodes: &[EpisodeInfo]) -> Result<Vec<Value>> {
        Err(crate::Error::model("predict_batch is not implemented"))
    }

    /// Whether this model defines the batched corner ([`predict_batch`](Self::predict_batch)).
    fn has_batch(&self) -> bool {
        false
    }

    /// 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`](Self::has_chunk_batch) is true.
    /// `execution_horizon` is the runtime's execution prefix (see
    /// [`predict_chunk`](Self::predict_chunk)); `episodes` is row-aligned with
    /// `inputs` (see [`predict_batch`](Self::predict_batch)). Default unimplemented
    /// (gated by the flag).
    fn predict_chunk_batch(
        &self,
        _inputs: Vec<Value>,
        _execution_horizon: u32,
        _episodes: &[EpisodeInfo],
    ) -> Result<Vec<Value>> {
        Err(crate::Error::model(
            "predict_chunk_batch is not implemented",
        ))
    }

    /// Whether this model defines the batched chunk corner ([`predict_chunk_batch`](Self::predict_chunk_batch)).
    fn has_chunk_batch(&self) -> bool {
        false
    }

    /// 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).
    fn predict_spec_less(&self, observation: ModelObservation) -> Result<Vec<SpaceValue>>;

    /// Spec-less route with a pinned execution horizon: like
    /// [`predict_spec_less`](Self::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 predict_spec_less_chunked(
        &self,
        observation: ModelObservation,
        _execution_horizon: u32,
    ) -> Result<super::handler::PredictFrames> {
        Ok(super::handler::PredictFrames {
            actions: self.predict_spec_less(observation)?,
            replay: Vec::new(),
        })
    }

    /// 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 allow_fusion(&self) -> bool {
        false
    }

    /// 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 native_chunk(&self) -> Option<u32> {
        None
    }

    /// 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_episode_end(&self, _episode_id: &str) -> Result<()> {
        Ok(())
    }

    /// Fires once at shutdown (structurally-discovered model hook, e.g. free
    /// GPU). The engine separately clears all frame buffers.
    fn on_close(&self) -> Result<()> {
        Ok(())
    }
}

/// The resolved per-route state the engine caches at `ResolveAdapter`: the
/// native adapter, the obs/action spaces, and the two host holes. Built by a
/// [`RouteResolver`] (which has the model spec); held beside an episode-keyed
/// `FrameBuffers` inside the engine.
pub struct RouteConfig {
    pub(crate) adapter: ResolvedAdapter,
    pub(crate) observation_space: SpaceSpec,
    pub(crate) action_space: SpaceSpec,
    pub(crate) customs: Box<dyn CustomTransform + Send + Sync>,
    pub(crate) encodings: Box<dyn EncodingTransform + Send + Sync>,
    /// Runtime-chosen execution horizon: how many actions of each predicted chunk the
    /// runtime executes before re-planning, set by the engine from the `ResolveAdapter`
    /// pin (1 = no chunking). Defaulted to 1 by [`new`](RouteConfig::new); the resolver
    /// builds the spec-derived config and the engine stamps the horizon on top, since
    /// it is a runtime decision, not part of the model spec.
    pub(crate) execution_horizon: u32,
    /// The runtime delivers every replayed step as a history row (negotiated at
    /// resolve): the engine ingests rows before the request's own observation
    /// and holds every frame to consecutive steps.
    pub(crate) delivers_history: bool,
}

impl RouteConfig {
    /// Assemble a route config from its resolved parts.
    ///
    /// `customs` fills custom-input holes (e.g. [`NoCustoms`](rlmesh_adapters::v1::NoCustoms)
    /// for a declarative route); `encodings` repacks custom rotation encodings
    /// (e.g. [`NoEncodings`](rlmesh_adapters::v1::NoEncodings) for a route with none).
    pub fn new(
        adapter: ResolvedAdapter,
        observation_space: SpaceSpec,
        action_space: SpaceSpec,
        customs: Box<dyn CustomTransform + Send + Sync>,
        encodings: Box<dyn EncodingTransform + Send + Sync>,
    ) -> Self {
        Self {
            adapter,
            observation_space,
            action_space,
            customs,
            encodings,
            // Spec-derived default; the engine overwrites it with the route's
            // runtime-pinned execution_horizon at ResolveAdapter.
            execution_horizon: 1,
            delivers_history: false,
        }
    }
}

/// Resolves a route's [`RouteConfig`] from its env contract at `ResolveAdapter`.
///
/// Returns `None` for a spec-less route (the env sent `NO_ADAPTER`), which the
/// engine then serves through [`PredictFn::predict_spec_less`]. Runs off the
/// predict-serialization lock (see [`ModelRouteSetup`](super::handler::ModelRouteSetup)),
/// so the binding may do blocking host work (resolving a spec) here.
#[async_trait]
pub trait RouteResolver: Send + Sync {
    /// Resolve `route_key`'s config from its `env_contract`, or `None` for a
    /// spec-less route. An error fails route configuration.
    async fn resolve(
        &self,
        route_key: &str,
        env_contract: &EnvContract,
    ) -> Result<Option<RouteConfig>>;
}