1use crate::{components::{Component, Transform}, hierarchy, reflect::Field};
4use glam::Vec3;
5use kvas_ecs::World;
6use serde::{Deserialize, Serialize};
7
8#[derive(Clone, Debug, Serialize, Deserialize)]
10pub struct Example {
11 pub input: Vec<f32>,
13 pub target: Vec<f32>,
15}
16
17#[derive(Clone, Debug, Serialize, Deserialize)]
19pub struct Network {
20 pub version: u32,
22 pub inputs: usize,
24 pub hidden: usize,
26 pub outputs: usize,
28 pub w1: Vec<f32>,
30 pub b1: Vec<f32>,
32 pub w2: Vec<f32>,
34 pub b2: Vec<f32>,
36}
37
38impl Network {
39 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 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 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 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 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 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#[derive(Clone, Serialize, Deserialize)]
120#[serde(default)]
121pub struct NeuralAgent {
122 pub enabled: bool,
124 pub target: String,
126 pub speed: f32,
128 pub input_range: f32,
130 pub stop_distance: f32,
132 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
142pub 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
166pub 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}