24539d2be1
With the world renderer through world/1, nothing in production calls int_hyper_synth::run_net, decode, to_rgb8 or the fast/reference convolution dispatch run_net used (wai-web included). They move into the test-only int_hyper_synth::legacy, beside the guard: run_net and apply_layer on the scalar kernels (legacy::reference), with no dispatch; decode and to_rgb8. int_hyper_synth::kernel_equivalence goes with the dispatch it tested: its convolution equivalences are int_ops' whm1_convolutions_equal_the_interpreter (M4), its pad-factor test int_ops' padding_beyond_the_pad_factor_runs_on_the_ interpreter (gate 5), and its overflow and malformed-geometry cases cannot reach a loaded graph. The callers that were not production move to the lowering: - wai_decode_bench: the int_hyper kernel cases, the g_s decoder case and the layer profile run whm1/1 graphs through int_ops, as the decoders do; - tests/int_kernels_geometry.rs: the int_hyper edge geometry runs a WHM1 layer through its lowering under every policy; strides past 32 bits and planes past the cap are refused, not panicked; - tests/int_hyper_synth.rs: the pretrained synthesis and hyper-decoder hashes are checked through the lowering alone; - int_mlicpp's and int_world's direct decodes take run_net, decode and to_rgb8 from legacy; docs in int_kernels and int_transformer. Co-Authored-By: Claude Opus 5.5 <[email protected]>
474 lines
24 KiB
Rust
474 lines
24 KiB
Rust
//! Integer-exact INTERACTIVE WORLD MODEL — the byte-exact renderer for
|
||
//! `wai.world.model`: every frame is generated by integer arithmetic, so its hash is
|
||
//! the same on every machine (a float next-frame network's would not be).
|
||
//!
|
||
//! An action-conditioned next-frame model, fully integer:
|
||
//! ```text
|
||
//! frame_0 = g_s(z_0); z_{t+1} = clip( T( z_t ‖ action_t ) ); frame_{t+1} = g_s(z_{t+1})
|
||
//! ```
|
||
//! Both the transition `T` and the synthesis `g_s` are [`IntModel`]s (the int_hyper
|
||
//! WHM1 format); the action is folded in as `A` constant input planes, `A` being `T`'s
|
||
//! input channels less the latent's. The renderer runs through the integer operator set
|
||
//! (`extensions/integer-operators.md`): `T` and `g_s` lowered by `world/1` (§13.5,
|
||
//! informative) and loaded under `acc64` with their bounds certificates ([`ops`]), so
|
||
//! given `(T, g_s, z_0, action trace)` every frame is byte-identical on every machine,
|
||
//! and a value that could take a layer outside its exact integer range is refused
|
||
//! (`check` ranges) rather than computed. Paired with [`crate::world_model`]'s signed,
|
||
//! replay-equivalent action trace, both the trace AND the generated frames
|
||
//! hash-converge: a generative interactive world you can put in a receipt.
|
||
//!
|
||
//! The renderer is informative: `wai.world.model`'s payload is its action trace, and no
|
||
//! payload text yet binds a trace to a renderer's set. Its refusals ([`Refused`]) use the
|
||
//! integer payloads' codes: `model_mismatch` for a set that breaks `world/1`'s rules, a row
|
||
//! without one value per action plane, or a network that does not fit the latent;
|
||
//! `geometry_mismatch` for a latent that is not `C·lat·lat` values; `field_range` for a side
|
||
//! of 0, a move whose plane value leaves `i64`, or a value outside a `check` range;
|
||
//! `too_large` for a latent past `2^26` values.
|
||
|
||
use crate::int_hyper_synth::{fnv1a, IntModel, Refused};
|
||
|
||
pub use ops::{load, World};
|
||
|
||
/// The reference action set (must match the export tool): stay, +x, −x, +y, −y.
|
||
/// `(dx, dy)` are folded into the transition as constant planes (scaled to `2^afrac`).
|
||
pub const ACTIONS: [(i64, i64); 5] = [(0, 0), (1, 0), (-1, 0), (0, 1), (0, -1)];
|
||
|
||
/// Number of action planes appended to the latent for the transition (dx, dy, btn).
|
||
pub const A: usize = 3;
|
||
|
||
/// A decoded interactive clip: one HWC `u8` RGB frame per (initial + per-action) step.
|
||
pub struct WorldClip {
|
||
pub frames: Vec<Vec<u8>>,
|
||
pub frame_px: usize,
|
||
}
|
||
|
||
impl WorldClip {
|
||
/// FNV-1a over one frame's RGB bytes — the per-frame portable hash.
|
||
pub fn frame_hash(&self, i: usize) -> u64 {
|
||
fnv1a(&self.frames[i])
|
||
}
|
||
/// The clip hash = FNV-1a over the concatenated per-frame hash bytes.
|
||
pub fn clip_hash(&self) -> u64 {
|
||
let mut all = Vec::with_capacity(self.frames.len() * 8);
|
||
for f in &self.frames {
|
||
all.extend_from_slice(&fnv1a(f).to_be_bytes());
|
||
}
|
||
fnv1a(&all)
|
||
}
|
||
}
|
||
|
||
/// Render a clip from rows of action plane values at `2^afrac`, one value per plane of
|
||
/// `T` (§13.5's driver): the set loaded ([`load`]), then [`World::decode_planes`]. A
|
||
/// merged multi-actor log maps to one row per op (the acting agent's planes filled,
|
||
/// the others 0), so participants co-driving one world render the same frames.
|
||
pub fn decode_planes(t_model: &IntModel, gs_model: &IntModel, z0: &[i64], rows: &[Vec<i64>], c: usize, lat: usize) -> Result<WorldClip, Refused> {
|
||
ops::decode_planes(t_model, gs_model, z0, rows, c, lat)
|
||
}
|
||
|
||
/// Render a clip from a single agent's moves `(dx, dy)` ([`World::decode_deltas`]): the
|
||
/// path of a merged, canonically ordered multi-actor log from
|
||
/// [`crate::world_model::replay`], which renders to one byte-exact frame sequence on every
|
||
/// machine, whatever order the ops arrived in.
|
||
pub fn decode_deltas(t_model: &IntModel, gs_model: &IntModel, z0: &[i64], deltas: &[(i64, i64)], c: usize, lat: usize) -> Result<WorldClip, Refused> {
|
||
ops::decode_deltas(t_model, gs_model, z0, deltas, c, lat)
|
||
}
|
||
|
||
/// Render a clip from indices into [`ACTIONS`] (past its end, the last).
|
||
pub fn decode(t_model: &IntModel, gs_model: &IntModel, z0: &[i64], actions: &[usize], c: usize, lat: usize) -> Result<WorldClip, Refused> {
|
||
ops::decode(t_model, gs_model, z0, actions, c, lat)
|
||
}
|
||
|
||
/// The direct renderer, as `wai.world.model`'s integer renderer ran before it rendered
|
||
/// through the integer operator set: `T` and `g_s` run layer by layer by
|
||
/// [`crate::int_hyper_synth::legacy::run_net`], with no check of the values they are given.
|
||
/// Test-only (feature `legacy_int`): the tests hold the renderer to it, and the two differ
|
||
/// only on the planned breaks of `extensions/integer-operators.md` §15.3.
|
||
#[cfg(any(test, feature = "legacy_int"))]
|
||
pub mod legacy {
|
||
use super::*;
|
||
use crate::int_hyper_synth::legacy::{decode as synth_decode, run_net};
|
||
|
||
/// One integer transition step from an explicit set of action **planes**:
|
||
/// `z_{t+1} = clip( T( z_t ‖ plane₀ ‖ … ‖ plane_{A-1} ) )`. Each `planes_q[k]` is a
|
||
/// constant value (already at `2^afrac`) broadcast to one `lat·lat` input plane, so
|
||
/// the number of planes (`A`) is whatever the model expects — `A=3` for a single
|
||
/// agent (dx, dy, btn), `A=4` for two agents (dxA, dyA, dxB, dyB), etc. The next
|
||
/// latent is clamped to `[-2^afrac, 2^afrac]` (the bound the model trained under).
|
||
/// The generalized core every higher-level step path calls.
|
||
pub fn step_planes(t_model: &IntModel, z: &[i64], planes_q: &[i64], c: usize, lat: usize) -> Vec<i64> {
|
||
let one = 1i64 << t_model.afrac;
|
||
let plane = lat * lat;
|
||
let na = planes_q.len();
|
||
let mut inp = Vec::with_capacity((c + na) * plane);
|
||
inp.extend_from_slice(z);
|
||
for &v in planes_q {
|
||
inp.extend(std::iter::repeat_n(v, plane));
|
||
}
|
||
let (mut zn, _, _, _) = run_net(t_model, &inp, c + na, lat, lat);
|
||
for v in &mut zn {
|
||
*v = (*v).clamp(-one, one);
|
||
}
|
||
zn
|
||
}
|
||
|
||
/// One integer transition step from raw `(dx, dy)` deltas (single agent, `A=3`).
|
||
pub fn step_delta(t_model: &IntModel, z: &[i64], dx: i64, dy: i64, c: usize, lat: usize) -> Vec<i64> {
|
||
let one = 1i64 << t_model.afrac;
|
||
step_planes(t_model, z, &[dx * one, dy * one, 0], c, lat)
|
||
}
|
||
|
||
/// One integer transition step from an action index into [`ACTIONS`].
|
||
pub fn step(t_model: &IntModel, z: &[i64], action: usize, c: usize, lat: usize) -> Vec<i64> {
|
||
let (dx, dy) = ACTIONS[action.min(ACTIONS.len() - 1)];
|
||
step_delta(t_model, z, dx, dy, c, lat)
|
||
}
|
||
|
||
/// Decode an interactive world clip from a list of `(dx, dy)` deltas: from `z0`,
|
||
/// render `g_s(z0)`, then for each delta advance the integer transition and render.
|
||
/// The convergence path — a **merged, canonically-ordered multi-actor** action log
|
||
/// (from [`crate::world_model::replay`]) decodes to one byte-exact frame sequence on
|
||
/// every machine, regardless of network arrival order. Pure i64.
|
||
pub fn decode_deltas(t_model: &IntModel, gs_model: &IntModel, z0: &[i64], deltas: &[(i64, i64)], c: usize, lat: usize) -> WorldClip {
|
||
let mut z = z0.to_vec();
|
||
let (f0, frame_px, _) = synth_decode(gs_model, &z, c, lat, lat);
|
||
let mut frames = Vec::with_capacity(deltas.len() + 1);
|
||
frames.push(f0);
|
||
for &(dx, dy) in deltas {
|
||
z = step_delta(t_model, &z, dx, dy, c, lat);
|
||
let (f, _, _) = synth_decode(gs_model, &z, c, lat, lat);
|
||
frames.push(f);
|
||
}
|
||
WorldClip { frames, frame_px }
|
||
}
|
||
|
||
/// Decode an interactive world clip from an explicit sequence of action **plane**
|
||
/// rows (each row `planes_q` at `2^afrac`, length = the model's `A`). The general
|
||
/// path — single- or multi-agent. A merged multi-actor log maps to one plane row per
|
||
/// op (the acting agent's slot filled, others zero), so two participants co-drive one
|
||
/// world and every peer renders the IDENTICAL frames. Pure i64.
|
||
pub fn decode_planes(t_model: &IntModel, gs_model: &IntModel, z0: &[i64], rows: &[Vec<i64>], c: usize, lat: usize) -> WorldClip {
|
||
let mut z = z0.to_vec();
|
||
let (f0, frame_px, _) = synth_decode(gs_model, &z, c, lat, lat);
|
||
let mut frames = Vec::with_capacity(rows.len() + 1);
|
||
frames.push(f0);
|
||
for row in rows {
|
||
z = step_planes(t_model, &z, row, c, lat);
|
||
let (f, _, _) = synth_decode(gs_model, &z, c, lat, lat);
|
||
frames.push(f);
|
||
}
|
||
WorldClip { frames, frame_px }
|
||
}
|
||
|
||
/// Decode an interactive world clip from action indices (single-actor convenience).
|
||
pub fn decode(t_model: &IntModel, gs_model: &IntModel, z0: &[i64], actions: &[usize], c: usize, lat: usize) -> WorldClip {
|
||
let deltas: Vec<(i64, i64)> = actions.iter().map(|&a| ACTIONS[a.min(ACTIONS.len() - 1)]).collect();
|
||
decode_deltas(t_model, gs_model, z0, &deltas, c, lat)
|
||
}
|
||
|
||
/// The direct renderer with the direct guard the other integer decoders ran
|
||
/// ([`crate::int_hyper_synth::legacy::run_net_checked`]): before each layer of `T` and
|
||
/// of `g_s`, the values present are checked against the bounds of that layer's
|
||
/// arithmetic, and a layer that could leave the exact integer range is not run. The
|
||
/// direct renderer itself never checked; this is what the `check` ranges of `world/1`
|
||
/// restate (`extensions/integer-operators.md` §13.2), and what the tests hold the
|
||
/// renderer through the operator set to. Test-only (`legacy_int`).
|
||
#[cfg(any(test, feature = "legacy_int"))]
|
||
pub mod checked {
|
||
use super::*;
|
||
use crate::int_hyper_synth::legacy::run_net_checked;
|
||
use crate::int_hyper_synth::legacy::to_rgb8;
|
||
|
||
/// Where the guard refused: the network (`"T"` or `"g_s"`), the step (0 for frame
|
||
/// 0, `k` for the `k`-th row's transition and the frame after it), the layer whose
|
||
/// check failed, and whether that check was about the values (`out_of_range`) or
|
||
/// about the layer itself.
|
||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||
pub struct Refusal {
|
||
pub net: &'static str,
|
||
pub step: usize,
|
||
pub layer: usize,
|
||
pub out_of_range: bool,
|
||
}
|
||
|
||
/// [`super::decode_planes`], with each network run by `run_net_checked`.
|
||
pub fn decode_planes(t_model: &IntModel, gs_model: &IntModel, z0: &[i64], rows: &[Vec<i64>], c: usize, lat: usize) -> Result<WorldClip, Refusal> {
|
||
let at = |net, step| move |r: crate::int_hyper_synth::NetRefusal| Refusal { net, step, layer: r.layer, out_of_range: r.out_of_range };
|
||
let render = |z: &[i64], step: usize| -> Result<(Vec<u8>, usize), Refusal> {
|
||
let (x, oc, oh, ow) = run_net_checked(gs_model, z, c, lat, lat).map_err(at("g_s", step))?;
|
||
Ok((to_rgb8(&x, oc, oh, ow, gs_model.afrac), oh))
|
||
};
|
||
let one = 1i64 << t_model.afrac;
|
||
let plane = lat * lat;
|
||
let mut z = z0.to_vec();
|
||
let (f0, frame_px) = render(&z, 0)?;
|
||
let mut frames = vec![f0];
|
||
for (k, row) in rows.iter().enumerate() {
|
||
let mut inp = z;
|
||
for &v in row {
|
||
inp.extend(std::iter::repeat_n(v, plane));
|
||
}
|
||
let (mut zn, ..) = run_net_checked(t_model, &inp, c + row.len(), lat, lat).map_err(at("T", k + 1))?;
|
||
for v in &mut zn {
|
||
*v = (*v).clamp(-one, one);
|
||
}
|
||
z = zn;
|
||
frames.push(render(&z, k + 1)?.0);
|
||
}
|
||
Ok(WorldClip { frames, frame_px })
|
||
}
|
||
}
|
||
}
|
||
|
||
/// The renderer through the integer operator set (`extensions/integer-operators.md`):
|
||
/// the transition `T` and the synthesis `g_s` lowered by `world/1` (§13.5, informative;
|
||
/// §13.2 for each network's layers), each loaded under `acc64` with its bounds
|
||
/// certificate ([`crate::int_ops`]).
|
||
pub mod ops {
|
||
use super::*;
|
||
use crate::int_hyper_synth::{refused, RefusalKind, Refused, MAX_TENSOR};
|
||
use wai_ops::{Code, Graph, Loaded, Refusal, Tensor};
|
||
|
||
/// A refusal of the operator set as the renderer reports it: a value outside a `check`
|
||
/// range is `field_range`; the other codes keep their names (integer-payloads §8).
|
||
fn kind(r: Refusal, what: &'static str) -> Refused {
|
||
let k = match r.code {
|
||
Code::FieldRange => RefusalKind::FieldRange,
|
||
Code::TooLarge => RefusalKind::TooLarge,
|
||
Code::GeometryMismatch => RefusalKind::GeometryMismatch,
|
||
Code::StreamExhausted => RefusalKind::StreamExhausted,
|
||
Code::StreamNoncanonical => RefusalKind::StreamNoncanonical,
|
||
Code::ModelMismatch => RefusalKind::ModelMismatch,
|
||
};
|
||
refused(k, what)
|
||
}
|
||
|
||
const SET: &str = "T and g_s do not lower by world/1 to integer-operator graphs that load under acc64";
|
||
|
||
/// Both networks in the form `wai-ops` lowers them from ([`crate::int_ops::whm1_net`]),
|
||
/// and `C` as the lowering takes it.
|
||
fn nets(t: &IntModel, g_s: &IntModel, c: usize) -> Result<(wai_ops::lower::Net, wai_ops::lower::Net, u32), Refused> {
|
||
let c = u32::try_from(c).map_err(|_| refused(RefusalKind::ModelMismatch, SET))?;
|
||
let tn = crate::int_ops::whm1_net(t).map_err(|r| kind(r, SET))?;
|
||
Ok((tn, crate::int_ops::whm1_net(g_s).map_err(|r| kind(r, SET))?, c))
|
||
}
|
||
|
||
/// The set's two graphs, `T` then `g_s` (with the RGB8 step), as `world/1` writes them
|
||
/// from the models this crate parsed, for a latent of `c` channels
|
||
/// ([`wai_ops::lower::world_net`]).
|
||
pub fn lower(t: &IntModel, g_s: &IntModel, c: usize) -> Result<(Graph, Graph), Refused> {
|
||
let (tn, gn, c) = nets(t, g_s, c)?;
|
||
let mut roles = wai_ops::lower::world_net(&tn, &gn, c).map_err(|r| kind(r, SET))?.roles.into_iter().map(|(_, g)| g);
|
||
match (roles.next(), roles.next()) {
|
||
(Some(tg), Some(gg)) => Ok((tg, gg)),
|
||
_ => Err(refused(RefusalKind::ModelMismatch, SET)),
|
||
}
|
||
}
|
||
|
||
/// The set as the renderer holds it: `T` and `g_s` lowered by `world/1` and loaded
|
||
/// under `acc64`, with the latent's channel count `C` and the number of action planes
|
||
/// `A`, which is `T`'s input channel count less `C` (§13.5).
|
||
pub struct World {
|
||
pub(super) t: Loaded,
|
||
pub(super) g_s: Loaded,
|
||
pub(super) c: usize,
|
||
pub(super) a: usize,
|
||
pub(super) afrac: u32,
|
||
}
|
||
|
||
impl World {
|
||
/// `T`'s graph, loaded.
|
||
pub fn t(&self) -> &Loaded {
|
||
&self.t
|
||
}
|
||
/// `g_s`'s graph (ending with the RGB8 step), loaded.
|
||
pub fn g_s(&self) -> &Loaded {
|
||
&self.g_s
|
||
}
|
||
/// `C`, the latent's channel count.
|
||
pub fn channels(&self) -> usize {
|
||
self.c
|
||
}
|
||
/// `A`, the action planes a row carries: `T`'s input channels less `C`.
|
||
pub fn planes(&self) -> usize {
|
||
self.a
|
||
}
|
||
/// The fixed-point scale both networks share.
|
||
pub fn afrac(&self) -> u32 {
|
||
self.afrac
|
||
}
|
||
|
||
/// A row carries one value per action plane (§13.5's load checks, `model_mismatch`).
|
||
fn row_ok(&self, row: &[i64]) -> Result<(), Refused> {
|
||
if row.len() != self.a {
|
||
return Err(refused(RefusalKind::ModelMismatch, "a row does not carry one value per action plane of T"));
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// The elements of a latent `[C, lat, lat]`: a side of 0 is `field_range`, and a
|
||
/// latent past [`MAX_TENSOR`] elements `too_large`.
|
||
fn latent_len(&self, lat: usize) -> Result<usize, Refused> {
|
||
if lat == 0 {
|
||
return Err(refused(RefusalKind::FieldRange, "the latent's side is zero"));
|
||
}
|
||
let n = self.c.checked_mul(lat).and_then(|v| v.checked_mul(lat)).filter(|&n| n <= MAX_TENSOR);
|
||
n.ok_or(refused(RefusalKind::TooLarge, "the latent exceeds 2^26 elements"))
|
||
}
|
||
|
||
/// The latent `z` as a tensor `[C, lat, lat]` ([`World::latent_len`]); `z` of
|
||
/// another length is `geometry_mismatch`.
|
||
fn latent(&self, z: &[i64], lat: usize) -> Result<Tensor, Refused> {
|
||
if z.len() != self.latent_len(lat)? {
|
||
return Err(refused(RefusalKind::GeometryMismatch, "the latent is not C·lat·lat values"));
|
||
}
|
||
// C and lat are within 2^26, so each fits u32.
|
||
Ok(Tensor::new(self.c as u32, lat as u32, lat as u32, z.to_vec()))
|
||
}
|
||
|
||
/// The frame a latent `[C, lat, lat]` gives, `(channels, h, w)`, planned before
|
||
/// any value is read (`extensions/integer-operators.md` §4): `T` planned on that
|
||
/// shape must give it back, so that every transition reads what the one before it
|
||
/// wrote; then `g_s` must plan on it. A network that does not (it reads another
|
||
/// channel count, a size passes a cap, or a layer leaves no output) does not fit the
|
||
/// latent: `model_mismatch`, as a set that does not fit a payload's shapes is.
|
||
pub fn plan(&self, lat: usize) -> Result<(usize, usize, usize), Refused> {
|
||
use wai_ops::plan::run_plan;
|
||
self.latent_len(lat)?;
|
||
let shape = [(self.c as u64, lat as u64, lat as u64)];
|
||
let t = self.t.graph();
|
||
let misfit = |what| move |_| refused(RefusalKind::ModelMismatch, what);
|
||
let hw = run_plan(t, &shape, &vec![0; self.a]).map_err(misfit("T does not plan on the latent's shape"))?;
|
||
let o = t.outputs[0] as usize;
|
||
if (t.tensors[o].channels as u64, hw[o].0, hw[o].1) != shape[0] {
|
||
return Err(refused(RefusalKind::ModelMismatch, "T does not give the latent's shape back"));
|
||
}
|
||
let g = self.g_s.graph();
|
||
let hw = run_plan(g, &shape, &[]).map_err(misfit("g_s does not plan on the latent's shape"))?;
|
||
let o = g.outputs[0] as usize;
|
||
Ok((g.tensors[o].channels as usize, hw[o].0 as usize, hw[o].1 as usize))
|
||
}
|
||
|
||
/// `T(z)` with `row`'s values as its scalars, then the clamp to `[−2^afrac,
|
||
/// 2^afrac]` that ends `T`'s graph. A value outside a `check` range is `field_range`.
|
||
fn run_t(&self, z: Tensor, row: &[i64]) -> Result<Tensor, Refused> {
|
||
let shape = (z.c, z.h, z.w);
|
||
let out = crate::int_ops::run1(&self.t, z, row).map_err(|r| kind(r, "a latent or plane value is outside a check range of T"))?;
|
||
if (out.c, out.h, out.w) != shape {
|
||
return Err(refused(RefusalKind::ModelMismatch, "T does not give the latent's shape back"));
|
||
}
|
||
Ok(out)
|
||
}
|
||
|
||
/// `g_s(z)`, which ends with the RGB8 step, as row-major HWC bytes, with the
|
||
/// frame's height and width. A value outside a `check` range is `field_range`.
|
||
fn run_g(&self, z: Tensor) -> Result<(Vec<u8>, usize, usize), Refused> {
|
||
let out = crate::int_ops::run1(&self.g_s, z, &[]).map_err(|r| kind(r, "a latent value is outside a check range of g_s"))?;
|
||
let (c, h, w) = (out.c as usize, out.h as usize, out.w as usize);
|
||
let plane = h * w;
|
||
let mut hwc = vec![0u8; c * plane];
|
||
for (k, &v) in out.data.iter().enumerate() {
|
||
hwc[(k % plane) * c + k / plane] = v as u8;
|
||
}
|
||
Ok((hwc, h, w))
|
||
}
|
||
|
||
/// One transition: `z` (`[C, lat, lat]`) becomes `T(z)`, with `row`'s values, one per
|
||
/// action plane at `2^afrac`, as `T`'s scalars (§13.5).
|
||
pub fn step(&self, z: &[i64], row: &[i64], lat: usize) -> Result<Vec<i64>, Refused> {
|
||
self.row_ok(row)?;
|
||
let x = self.latent(z, lat)?;
|
||
Ok(self.run_t(x, row)?.data)
|
||
}
|
||
|
||
/// [`World::step`] for a single agent's move `(dx, dy)` ([`World::move_row`]).
|
||
pub fn step_delta(&self, z: &[i64], dx: i64, dy: i64, lat: usize) -> Result<Vec<i64>, Refused> {
|
||
let row = self.move_row(dx, dy)?;
|
||
self.step(z, &row, lat)
|
||
}
|
||
|
||
/// [`World::step_delta`] for an index into [`ACTIONS`] (past its end, the last).
|
||
pub fn step_action(&self, z: &[i64], action: usize, lat: usize) -> Result<Vec<i64>, Refused> {
|
||
let (dx, dy) = ACTIONS[action.min(ACTIONS.len() - 1)];
|
||
self.step_delta(z, dx, dy, lat)
|
||
}
|
||
|
||
/// One frame: `g_s(z)` as row-major HWC bytes, with its height and width.
|
||
pub fn render(&self, z: &[i64], lat: usize) -> Result<(Vec<u8>, usize, usize), Refused> {
|
||
let x = self.latent(z, lat)?;
|
||
self.run_g(x)
|
||
}
|
||
|
||
/// The row of a single agent's move `(dx, dy)` (§13.5): `dequantize(dx, afrac, 0)`,
|
||
/// `dequantize(dy, afrac, 0)` and 0. A set whose `A` is not 3 takes no such row
|
||
/// (`model_mismatch`), and a move whose value leaves `i64` is `field_range`.
|
||
pub fn move_row(&self, dx: i64, dy: i64) -> Result<Vec<i64>, Refused> {
|
||
if self.a != 3 {
|
||
return Err(refused(RefusalKind::ModelMismatch, "a single agent's move is a row of three planes, and T takes another count"));
|
||
}
|
||
let q = |d: i64| wai_ops::stage::dequantize(d, self.afrac, 0).map_err(|_| refused(RefusalKind::FieldRange, "a move's plane value leaves i64"));
|
||
Ok(vec![q(dx)?, q(dy)?, 0])
|
||
}
|
||
|
||
/// A clip (§13.5's driver): frame 0 is `g_s(z_0)`; for each row in order, `z`
|
||
/// becomes `T(z)` with the row's values as the scalars, and the next frame is
|
||
/// `g_s(z)`. Before the first frame: every row carries `A` values
|
||
/// (`model_mismatch`), `z_0` is `[C, lat, lat]`, and both graphs plan on it
|
||
/// ([`World::plan`]).
|
||
pub fn decode_planes(&self, z0: &[i64], rows: &[Vec<i64>], lat: usize) -> Result<WorldClip, Refused> {
|
||
for row in rows {
|
||
self.row_ok(row)?;
|
||
}
|
||
let mut z = self.latent(z0, lat)?;
|
||
self.plan(lat)?;
|
||
let (mut frames, mut frame_px) = (Vec::with_capacity(rows.len() + 1), 0);
|
||
for k in 0..=rows.len() {
|
||
if k > 0 {
|
||
z = self.run_t(z, &rows[k - 1])?;
|
||
}
|
||
// The last frame takes the latent, which no transition reads after it.
|
||
let x = if k == rows.len() { std::mem::replace(&mut z, Tensor::new(0, 0, 0, Vec::new())) } else { z.clone() };
|
||
let (f, h, _) = self.run_g(x)?;
|
||
if k == 0 {
|
||
frame_px = h;
|
||
}
|
||
frames.push(f);
|
||
}
|
||
Ok(WorldClip { frames, frame_px })
|
||
}
|
||
|
||
/// [`World::decode_planes`] for a single agent's moves, each a row
|
||
/// ([`World::move_row`]).
|
||
pub fn decode_deltas(&self, z0: &[i64], deltas: &[(i64, i64)], lat: usize) -> Result<WorldClip, Refused> {
|
||
let rows = deltas.iter().map(|&(dx, dy)| self.move_row(dx, dy)).collect::<Result<Vec<_>, _>>()?;
|
||
self.decode_planes(z0, &rows, lat)
|
||
}
|
||
}
|
||
|
||
/// [`World::decode_planes`] of the set `T`, `g_s` for a latent of `c` channels, loaded
|
||
/// first ([`load`]).
|
||
pub fn decode_planes(t: &IntModel, g_s: &IntModel, z0: &[i64], rows: &[Vec<i64>], c: usize, lat: usize) -> Result<WorldClip, Refused> {
|
||
load(t, g_s, c)?.decode_planes(z0, rows, lat)
|
||
}
|
||
|
||
/// [`World::decode_deltas`] of the set, loaded first ([`load`]).
|
||
pub fn decode_deltas(t: &IntModel, g_s: &IntModel, z0: &[i64], deltas: &[(i64, i64)], c: usize, lat: usize) -> Result<WorldClip, Refused> {
|
||
load(t, g_s, c)?.decode_deltas(z0, deltas, lat)
|
||
}
|
||
|
||
/// [`decode_deltas`] for indices into [`ACTIONS`] (past its end, the last).
|
||
pub fn decode(t: &IntModel, g_s: &IntModel, z0: &[i64], actions: &[usize], c: usize, lat: usize) -> Result<WorldClip, Refused> {
|
||
let deltas: Vec<(i64, i64)> = actions.iter().map(|&a| ACTIONS[a.min(ACTIONS.len() - 1)]).collect();
|
||
decode_deltas(t, g_s, z0, &deltas, c, lat)
|
||
}
|
||
|
||
/// The load checks of `world/1` (§13.5, `model_mismatch`): both networks' `afrac` in
|
||
/// `1..=30` and equal, `A ≥ 1`, and each network's rules (§13.2); then `T`, then `g_s`,
|
||
/// lowered and loaded under `acc64` ([`wai_ops::lower::world_net_loaded`]: the result of
|
||
/// loading the lowered graphs' bytes, without writing and parsing them again).
|
||
pub fn load(t: &IntModel, g_s: &IntModel, c: usize) -> Result<World, Refused> {
|
||
let (tn, gn, c32) = nets(t, g_s, c)?;
|
||
let (tl, gl) = wai_ops::lower::world_net_loaded(tn, gn, c32, crate::int_ops::PROFILE).map_err(|r| kind(r, SET))?;
|
||
let a = tl.graph().scalars.len();
|
||
Ok(World { t: tl, g_s: gl, c, a, afrac: t.afrac })
|
||
}
|
||
}
|