infotheory/aixi/
planner_runtime.rs

1//! Compiled-spec bindings for planner-run documents.
2//!
3//! This module keeps filesystem asset loading and environment construction at
4//! the edge of the planner runtime. Controller execution lives in
5//! [`crate::aixi::planner_agent`].
6
7use crate::aixi::common::ActionAlphabet;
8use crate::aixi::environment::Environment;
9#[cfg(feature = "aixi-gameengine")]
10use crate::aixi::gameengine::build_builtin_environment as build_gameengine_builtin_environment;
11#[cfg(feature = "vm")]
12use crate::aixi::vm_nyx::{NyxVmConfig, NyxVmEnvironment};
13use crate::aixi::warmstart::{
14    WarmStartExactJhTeacherDataset, validate_warmstart_teacher_dataset_for_compiled_planner_run,
15};
16use crate::spec::{self, AssetRef, BuiltinEnvironmentSpec, CompiledPlannerRunSpec, SpecDocument};
17use std::path::Path;
18
19/// Load and validate a warm-start teacher dataset asset referenced by a compiled planner run.
20pub fn load_warmstart_exact_jh_teacher_dataset(
21    compiled: &CompiledPlannerRunSpec,
22    asset_id: &str,
23) -> anyhow::Result<WarmStartExactJhTeacherDataset> {
24    let binding = compiled
25        .resolved_assets()
26        .iter()
27        .find(|entry| entry.id == asset_id)
28        .ok_or_else(|| anyhow::anyhow!("unknown warm-start teacher_dataset_asset '{asset_id}'"))?;
29    let AssetRef::Filesystem(path) = &binding.asset;
30    let bytes = std::fs::read(path).map_err(|err| {
31        anyhow::anyhow!(
32            "failed to read warm-start teacher_dataset_asset '{}': {err}",
33            path.display()
34        )
35    })?;
36    let teacher =
37        WarmStartExactJhTeacherDataset::from_json_slice(&bytes).map_err(anyhow::Error::msg)?;
38    validate_warmstart_teacher_dataset_for_compiled_planner_run(compiled, &teacher)
39        .map_err(anyhow::Error::msg)?;
40    Ok(teacher)
41}
42
43/// Validate a teacher dataset contract against a compiled planner run.
44pub fn validate_warmstart_exact_jh_teacher_contract(
45    compiled: &CompiledPlannerRunSpec,
46    teacher: &WarmStartExactJhTeacherDataset,
47) -> anyhow::Result<()> {
48    crate::aixi::warmstart::validate_warmstart_teacher_against_compiled_planner_run(
49        compiled,
50        &teacher.contract,
51    )
52    .map_err(|err| anyhow::anyhow!("{err}"))
53}
54
55/// Compile a planner_run document from a filesystem path.
56pub fn compile_planner_run_document(
57    path: &str,
58    caller: &str,
59) -> anyhow::Result<CompiledPlannerRunSpec> {
60    let config_dir = Path::new(path).parent().unwrap_or(Path::new("."));
61    let document = spec::load_spec_document(path).map_err(anyhow::Error::msg)?;
62    match document {
63        SpecDocument::PlannerRun(spec) => spec
64            .compile_in(&spec::SpecEnvironment::new(config_dir))
65            .map_err(anyhow::Error::msg),
66        other => Err(anyhow::anyhow!(
67            "{caller} expects a planner_run document, found kind '{}'",
68            other.kind_str()
69        )),
70    }
71}
72
73fn build_builtin_environment(spec: BuiltinEnvironmentSpec) -> anyhow::Result<Box<dyn Environment>> {
74    #[cfg(feature = "aixi-gameengine")]
75    {
76        build_gameengine_builtin_environment(spec).map_err(anyhow::Error::new)
77    }
78    #[cfg(not(feature = "aixi-gameengine"))]
79    {
80        Err(anyhow::anyhow!(
81            "builtin environment '{}' requires feature 'aixi-gameengine'",
82            spec.canonical_name()
83        ))
84    }
85}
86
87/// Build the environment declared by a compiled planner run.
88pub fn build_planner_environment(
89    compiled: &CompiledPlannerRunSpec,
90) -> anyhow::Result<(Box<dyn Environment>, &'static str)> {
91    #[allow(unreachable_patterns)]
92    match &compiled.canonical_spec().environment {
93        spec::EnvironmentSpec::Builtin { builtin } => Ok((
94            build_builtin_environment(*builtin)?,
95            builtin.canonical_name(),
96        )),
97        #[cfg(feature = "vm")]
98        spec::EnvironmentSpec::NyxVm(vm) => {
99            let config = NyxVmConfig::from_environment_spec(vm, compiled.resolved_assets())
100                .map_err(anyhow::Error::msg)?;
101            Ok((Box::new(NyxVmEnvironment::new(config)?), "vm"))
102        }
103        #[cfg(not(feature = "vm"))]
104        other => Err(anyhow::anyhow!(
105            "unsupported environment variant '{}' in this build",
106            other.kind_str()
107        )),
108    }
109}
110
111/// Validate that the environment action alphabet matches the planner interface.
112pub fn validate_action_alphabet(
113    compiled: &CompiledPlannerRunSpec,
114    env: &dyn Environment,
115) -> anyhow::Result<()> {
116    let actual: ActionAlphabet = env.get_num_actions();
117    let expected: ActionAlphabet = compiled.interface().agent_actions;
118    if actual != expected {
119        return Err(anyhow::anyhow!(
120            "action_alphabet_mismatch: planner interface declares {} actions but environment exposes {}",
121            expected,
122            actual
123        ));
124    }
125    Ok(())
126}
127
128/// Validate the full environment interface against the compiled planner contract.
129pub(crate) fn validate_environment_interface(
130    compiled: &CompiledPlannerRunSpec,
131    env: &dyn Environment,
132) -> anyhow::Result<()> {
133    validate_action_alphabet(compiled, env)?;
134    let interface = compiled.interface();
135    let expected_action_bits: usize = interface.agent_actions.action_bits();
136    let actual_action_bits: usize = env.get_action_bits();
137    if actual_action_bits != expected_action_bits {
138        return Err(anyhow::anyhow!(
139            "action_bits_mismatch: planner interface declares {} action bits for {} actions but environment exposes {}",
140            expected_action_bits,
141            interface.agent_actions,
142            actual_action_bits
143        ));
144    }
145    let actual_observation_bits: usize = env.get_observation_bits();
146    if actual_observation_bits != interface.observation_bits {
147        return Err(anyhow::anyhow!(
148            "observation_bits_mismatch: planner interface declares {} observation bits but environment exposes {}",
149            interface.observation_bits,
150            actual_observation_bits
151        ));
152    }
153    let actual_reward_bits: usize = env.get_reward_bits();
154    if actual_reward_bits != interface.reward_bits {
155        return Err(anyhow::anyhow!(
156            "reward_bits_mismatch: planner interface declares {} reward bits but environment exposes {}",
157            interface.reward_bits,
158            actual_reward_bits
159        ));
160    }
161    Ok(())
162}
infotheory/aixi/planner_runtime.rs

infotheory/aixi/
planner_runtime.rs
