Skip to main content

kvas_engine/
neural.rs

1//! Small CPU supervised networks with portable JSON weights and a planar steering component.
2//! This is imitation/regression training, not reinforcement learning, navigation or NVIDIA inference.
3use crate::{components::{Component, Transform}, hierarchy, reflect::Field};
4use glam::Vec3;
5use kvas_ecs::World;
6use serde::{Deserialize, Serialize};
7
8/// One supervised training pair. Normalize features and targets before training.
9#[derive(Clone, Debug, Serialize, Deserialize)]
10pub struct Example {
11    /// Input features.
12    pub input: Vec<f32>,
13    /// Desired outputs.
14    pub target: Vec<f32>,
15}
16
17/// One tanh hidden layer and a linear output layer, stored as row-major JSON weights.
18#[derive(Clone, Debug, Serialize, Deserialize)]
19pub struct Network {
20    /// Format version; currently 1.
21    pub version: u32,
22    /// Number of input features.
23    pub inputs: usize,
24    /// Number of tanh neurons.
25    pub hidden: usize,
26    /// Number of predicted outputs.
27    pub outputs: usize,
28    /// Hidden-layer weights, hidden × inputs.
29    pub w1: Vec<f32>,
30    /// Hidden-layer biases.
31    pub b1: Vec<f32>,
32    /// Output-layer weights, outputs × hidden.
33    pub w2: Vec<f32>,
34    /// Output-layer biases.
35    pub b2: Vec<f32>,
36}
37
38impl Network {
39    /// Deterministic Xavier-like initialization. Each dimension must be 1..=128.
40    pub fn new(inputs: usize, hidden: usize, outputs: usize, seed: u64) -> Result<Self, String> {
41        if [inputs, hidden, outputs].iter().any(|&n| n == 0 || n > 128) { return Err("network dimensions must be 1..128".into()); }
42        let mut state = seed;
43        let mut weight = |scale: f32| {
44            state = state.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
45            (((state >> 40) as f32 / 16777215.0) * 2.0 - 1.0) * scale
46        };
47        Ok(Self { version: 1, inputs, hidden, outputs, w1: (0..hidden * inputs).map(|_| weight((6.0 / (inputs + hidden) as f32).sqrt())).collect(), b1: vec![0.0; hidden], w2: (0..outputs * hidden).map(|_| weight((6.0 / (hidden + outputs) as f32).sqrt())).collect(), b2: vec![0.0; outputs] })
48    }
49
50    /// Rejects unsupported versions, oversized dimensions, corrupt shapes and non-finite weights.
51    pub fn validate(&self) -> Result<(), String> {
52        if self.version != 1 || [self.inputs, self.hidden, self.outputs].iter().any(|&n| n == 0 || n > 128) { return Err("unsupported network version or dimensions".into()); }
53        if self.w1.len() != self.inputs * self.hidden || self.b1.len() != self.hidden || self.w2.len() != self.hidden * self.outputs || self.b2.len() != self.outputs { return Err("network weight shape mismatch".into()); }
54        if self.w1.iter().chain(&self.b1).chain(&self.w2).chain(&self.b2).any(|v| !v.is_finite()) { return Err("non-finite network weights".into()); }
55        Ok(())
56    }
57
58    /// Parses and validates portable weights.
59    pub fn from_json(json: &str) -> Result<Self, String> {
60        let network: Self = serde_json::from_str(json).map_err(|e| e.to_string())?;
61        network.validate()?; Ok(network)
62    }
63
64    /// Atomically saves validated weights. Parent directory must exist.
65    pub fn save(&self, path: &std::path::Path) -> Result<(), String> {
66        self.validate()?;
67        let bytes = serde_json::to_vec_pretty(self).map_err(|e| e.to_string())?;
68        crate::file_io::atomic_write(path, &bytes).map_err(|e| e.to_string())
69    }
70
71    fn forward(&self, input: &[f32]) -> (Vec<f32>, Vec<f32>) {
72        let hidden: Vec<f32> = (0..self.hidden).map(|h| (self.b1[h] + (0..self.inputs).map(|i| input[i] * self.w1[h * self.inputs + i]).sum::<f32>()).tanh()).collect();
73        let output = (0..self.outputs).map(|o| self.b2[o] + (0..self.hidden).map(|h| hidden[h] * self.w2[o * self.hidden + h]).sum::<f32>()).collect();
74        (hidden, output)
75    }
76
77    /// Predicts finite outputs or reports invalid features/weights.
78    pub fn predict(&self, input: &[f32]) -> Result<Vec<f32>, String> {
79        self.validate()?;
80        if input.len() != self.inputs || input.iter().any(|v| !v.is_finite()) { return Err("invalid network input".into()); }
81        let output = self.forward(input).1;
82        if output.iter().any(|v| !v.is_finite()) { return Err("network output overflow".into()); }
83        Ok(output)
84    }
85
86    /// One deterministic SGD epoch, returning pre-update mean squared error.
87    /// Gradient clipping bounds each parameter update. Call repeatedly on a worker, with cancellation.
88    pub fn train_epoch(&mut self, examples: &[Example], rate: f32) -> Result<f32, String> {
89        self.validate()?;
90        if !rate.is_finite() || rate <= 0.0 || rate > 1.0 || examples.is_empty() || examples.len() > 100_000 { return Err("invalid learning rate or dataset size".into()); }
91        for e in examples {
92            if e.input.len() != self.inputs || e.target.len() != self.outputs || e.input.iter().chain(&e.target).any(|v| !v.is_finite() || v.abs() > 1000.0) { return Err("training data shape mismatch or invalid feature; normalize data".into()); }
93        }
94        let mut mse = 0.0;
95        for e in examples {
96            let (h, y) = self.forward(&e.input);
97            let errors: Vec<_> = y.iter().zip(&e.target).map(|(a, b)| a - b).collect();
98            mse += errors.iter().map(|v| v * v).sum::<f32>() / self.outputs as f32;
99            let grad_out: Vec<_> = errors.iter().map(|v| (2.0 * v / self.outputs as f32).clamp(-10.0, 10.0)).collect();
100            let grad_hidden: Vec<_> = (0..self.hidden).map(|j| ((0..self.outputs).map(|o| grad_out[o] * self.w2[o * self.hidden + j]).sum::<f32>() * (1.0 - h[j] * h[j])).clamp(-10.0, 10.0)).collect();
101            for o in 0..self.outputs {
102                self.b2[o] -= rate * grad_out[o];
103                for j in 0..self.hidden { self.w2[o * self.hidden + j] -= rate * (grad_out[o] * h[j]).clamp(-10.0, 10.0); }
104            }
105            for j in 0..self.hidden {
106                self.b1[j] -= rate * grad_hidden[j];
107                for i in 0..self.inputs { self.w1[j * self.inputs + i] -= rate * (grad_hidden[j] * e.input[i]).clamp(-10.0, 10.0); }
108            }
109        }
110        self.validate()?;
111        let result = mse / examples.len() as f32;
112        if !result.is_finite() { return Err("training diverged".into()); }
113        Ok(result)
114    }
115}
116
117/// Small supervised steering controller. Inputs are normalized world X/Z offsets to a named target.
118/// Outputs are planar movement X/Z. Does not perform pathfinding or control dynamic rigid bodies.
119#[derive(Clone, Serialize, Deserialize)]
120#[serde(default)]
121pub struct NeuralAgent {
122    /// Enable CPU inference during Play.
123    pub enabled: bool,
124    /// Exact scene name of the target.
125    pub target: String,
126    /// Maximum world movement speed in metres per second.
127    pub speed: f32,
128    /// Distance represented by a feature magnitude of one.
129    pub input_range: f32,
130    /// Stop within this world-space distance.
131    pub stop_distance: f32,
132    /// Embedded trained weights, saved with the scene.
133    pub network: Network,
134}
135impl Default for NeuralAgent {
136    fn default() -> Self { Self { enabled: false, target: String::new(), speed: 2.0, input_range: 10.0, stop_distance: 0.3, network: Network::new(2, 8, 2, 42).unwrap() } }
137}
138impl Component for NeuralAgent {
139    fn fields(&mut self) -> Vec<Field<'_>> { vec![Field::bool("enabled", &mut self.enabled), Field::entity("target", &mut self.target), Field::f32("speed", &mut self.speed), Field::f32("input_range", &mut self.input_range), Field::f32("stop_distance", &mut self.stop_distance)] }
140}
141
142/// Advances enabled, valid steering agents. Invalid models/settings leave entities untouched.
143pub fn update(world: &mut World, dt: f32) {
144    if !dt.is_finite() || dt <= 0.0 { return; }
145    for e in world.entities_with::<NeuralAgent>() {
146        let Some(agent) = world.get::<NeuralAgent>(e) else { continue };
147        if !agent.enabled || !agent.speed.is_finite() || agent.speed < 0.0 || !agent.input_range.is_finite() || agent.input_range <= 0.0 || !agent.stop_distance.is_finite() || agent.stop_distance < 0.0 || agent.network.inputs != 2 || agent.network.outputs != 2 { continue; }
148        if world.has::<crate::components::RigidBody>(e) { continue; }
149        let Some(target) = hierarchy::find_by_name(world, &agent.target) else { continue };
150        let position = hierarchy::global_matrix(world, e).w_axis.truncate();
151        let delta = hierarchy::global_matrix(world, target).w_axis.truncate() - position;
152        let planar = glam::Vec2::new(delta.x, delta.z);
153        let distance = planar.length();
154        if distance <= agent.stop_distance { continue; }
155        let input = (planar / agent.input_range).clamp(glam::Vec2::NEG_ONE, glam::Vec2::ONE);
156        let Ok(output) = agent.network.predict(&input.to_array()) else { continue };
157        let direction = glam::Vec2::new(output[0], output[1]).clamp_length_max(1.0);
158        let step = (agent.speed * dt).min((distance - agent.stop_distance).max(0.0));
159        let desired = position + Vec3::new(direction.x, 0.0, direction.y) * step;
160        let parent = hierarchy::parent_of(world, e).map_or(glam::Mat4::IDENTITY, |p| hierarchy::global_matrix(world, p));
161        let local = parent.inverse().transform_point3(desired);
162        if local.is_finite() { if let Some(pose) = world.get_mut::<Transform>(e) { pose.position = local; } }
163    }
164}
165
166/// Reproducible steering examples: normalized X/Z offsets mapped to desired X/Z movement.
167pub fn steering_examples() -> Vec<Example> {
168    [-1.0f32, -0.5, 0.0, 0.5, 1.0].into_iter().flat_map(|x| [-1.0f32, -0.5, 0.0, 0.5, 1.0].into_iter().map(move |z| Example { input: vec![x, z], target: vec![x, z] })).collect()
169}
170
171#[cfg(test)]
172mod tests {
173    use super::*;
174    #[test]
175    fn training_reduces_error_and_weights_round_trip() {
176        let mut model = Network::new(2, 8, 2, 42).unwrap(); let data = steering_examples();
177        let initial = model.train_epoch(&data, 0.02).unwrap();
178        let mut loss = initial;
179        for _ in 0..500 { loss = model.train_epoch(&data, 0.02).unwrap(); }
180        assert!(loss < initial * 0.05, "{initial} -> {loss}");
181        let output = model.predict(&[0.75, -0.75]).unwrap(); assert!(output[0] > 0.6 && output[1] < -0.6);
182        let loaded = Network::from_json(&serde_json::to_string(&model).unwrap()).unwrap();
183        assert_eq!(output, loaded.predict(&[0.75, -0.75]).unwrap());
184        model.w1.clear(); assert!(model.predict(&[0.0, 0.0]).is_err());
185    }
186}