separate depth from inputs

src/inference.rs

@@ -32,10 +32,9 @@ impl SimulateSubcommand {
 				.to_str()
 				.unwrap()
 				.to_owned();
-			file_name.push_str("_replay");
+			file_name.push_str("_replay.snfb");
 			let mut path = self.model_file.clone();
 			path.set_file_name(file_name);
-			path.set_extension("snfb");
 			path
 		});
 		inference(
@@ -50,12 +49,12 @@ impl SimulateSubcommand {
 use burn::prelude::*;
 
 use crate::inputs::InputGenerator;
-use crate::net::{INPUT, InferenceBackend, Net};
+use crate::net::{DEPTH_SIZE, InferenceBackend, Net, POSITION_HISTORY_SIZE};
 
 use strafesnet_common::instruction::TimedInstruction;
 use strafesnet_common::mouse::MouseState;
 use strafesnet_common::physics::{
-	Instruction as PhysicsInputInstruction, MiscInstruction, ModeInstruction, MouseInstruction,
+	Instruction as PhysicsInputInstruction, ModeInstruction, MouseInstruction,
 	SetControlInstruction, Time as PhysicsTime,
 };
 use strafesnet_physics::physics::{PhysicsContext, PhysicsData, PhysicsState};
@@ -172,6 +171,7 @@ fn inference(
 
 	const STEP: PhysicsTime = PhysicsTime::from_millis(10);
 	let mut input_floats = Vec::new();
+	let mut depth_floats = Vec::new();
 	// setup agent-simulation feedback loop
 	for _ in 0..20 * 100 {
 		// generate inputs
@@ -180,14 +180,23 @@ fn inference(
 			frame_state.pos(time) - start_offset,
 			frame_state.angles(),
 			&mut input_floats,
+			&mut depth_floats,
 		);
 
 		// inference
 		let inputs = Tensor::from_data(
-			TensorData::new(input_floats.clone(), Shape::new([1, INPUT])),
+			TensorData::new(input_floats.clone(), Shape::new([1, POSITION_HISTORY_SIZE])),
 			&device,
 		);
-		let outputs = model.forward(inputs).into_data().into_vec::<f32>().unwrap();
+		let depth = Tensor::from_data(
+			TensorData::new(depth_floats.clone(), Shape::new([1, DEPTH_SIZE])),
+			&device,
+		);
+		let outputs = model
+			.forward(inputs, depth)
+			.into_data()
+			.into_vec::<f32>()
+			.unwrap();
 
 		let &[
 			move_forward,
@@ -231,6 +240,7 @@ fn inference(
 		time = next_time;
 
 		// clear
+		depth_floats.clear();
 		input_floats.clear();
 	}
 

src/inputs.rs

@@ -14,7 +14,7 @@ pub struct InputGenerator {
 	graphics_texture_view: wgpu::TextureView,
 	output_staging_buffer: wgpu::Buffer,
 	texture_data: Vec<u8>,
-	position_history: Vec<glam::Vec3>,
+	position_history: Vec<(glam::Vec3, glam::Vec2)>,
 }
 impl InputGenerator {
 	pub fn new(map: &strafesnet_common::map::CompleteMap) -> Self {
@@ -77,26 +77,34 @@ impl InputGenerator {
 			position_history,
 		}
 	}
-	pub fn generate_inputs(&mut self, pos: glam::Vec3, angles: glam::Vec2, inputs: &mut Vec<f32>) {
+	pub fn generate_inputs(
+		&mut self,
+		pos: glam::Vec3,
+		angles: glam::Vec2,
+		inputs: &mut Vec<f32>,
+		depth: &mut Vec<f32>,
+	) {
 		// write position history to model inputs
 		if !self.position_history.is_empty() {
 			let camera = strafesnet_graphics::graphics::view_inv(pos, angles).inverse();
-			for &pos in self.position_history.iter().rev() {
+			for &(pos, ang) in self.position_history.iter().rev() {
 				let relative_pos = camera.transform_vector3(pos);
+				let relative_ang = glam::vec2(angles.x - ang.x, ang.y);
 				inputs.extend_from_slice(&relative_pos.to_array());
+				inputs.extend_from_slice(&relative_ang.to_array());
 			}
 		}
 		// fill remaining history with zeroes
 		for _ in self.position_history.len()..POSITION_HISTORY {
-			inputs.extend_from_slice(&[0.0, 0.0, 0.0]);
+			inputs.extend_from_slice(&[0.0, 0.0, 0.0, 0.0, 0.0]);
 		}
 
 		// track position history
 		if self.position_history.len() < POSITION_HISTORY {
-			self.position_history.push(pos);
+			self.position_history.push((pos, angles));
 		} else {
 			self.position_history.rotate_left(1);
-			*self.position_history.last_mut().unwrap() = pos;
+			*self.position_history.last_mut().unwrap() = (pos, angles);
 		}
 
 		let mut encoder = self
@@ -153,7 +161,7 @@ impl InputGenerator {
 
 		// discombolulate stride
 		for y in 0..SIZE.y {
-			inputs.extend(
+			depth.extend(
 				self.texture_data[(STRIDE_SIZE * y) as usize
 					..(STRIDE_SIZE * y + SIZE.x * size_of::<f32>() as u32) as usize]
 					.chunks_exact(4)

src/net.rs

@@ -6,8 +6,10 @@ pub type InferenceBackend = burn::backend::Cuda<f32>;
 pub type TrainingBackend = Autodiff<InferenceBackend>;
 
 pub const SIZE: glam::UVec2 = glam::uvec2(64, 36);
+pub const DEPTH_SIZE: usize = (SIZE.x * SIZE.y) as usize;
 pub const POSITION_HISTORY: usize = 10;
-pub const INPUT: usize = (SIZE.x * SIZE.y) as usize + POSITION_HISTORY * 3;
+pub const POSITION_HISTORY_SIZE: usize = POSITION_HISTORY * 5;
+const INPUT: usize = DEPTH_SIZE + POSITION_HISTORY_SIZE;
 pub const HIDDEN: [usize; 3] = [INPUT >> 3, INPUT >> 5, INPUT >> 7];
 // MoveForward
 // MoveLeft
@@ -47,9 +49,10 @@ impl<B: Backend> Net<B> {
 			activation: Relu::new(),
 		}
 	}
-	pub fn forward(&self, input: Tensor<B, 2>) -> Tensor<B, 2> {
-		let x = self.input.forward(input);
-		let x = self.dropout.forward(x);
+	pub fn forward(&self, input: Tensor<B, 2>, depth: Tensor<B, 2>) -> Tensor<B, 2> {
+		let x = self.dropout.forward(depth);
+		let x = Tensor::cat(vec![input, x], 1);
+		let x = self.input.forward(x);
 		let mut x = self.activation.forward(x);
 		for layer in &self.hidden {
 			x = layer.forward(x);

src/training.rs

@@ -40,7 +40,7 @@ use burn::optim::{AdamConfig, GradientsParams, Optimizer};
 use burn::prelude::*;
 
 use crate::inputs::InputGenerator;
-use crate::net::{INPUT, Net, OUTPUT, TrainingBackend};
+use crate::net::{DEPTH_SIZE, Net, OUTPUT, POSITION_HISTORY_SIZE, TrainingBackend};
 
 use strafesnet_roblox_bot_file::v0;
 
@@ -72,8 +72,10 @@ fn training(
 	// training data
 	let training_samples = timelines.input_events.len() - 1;
 
-	let input_size = INPUT * size_of::<f32>();
+	let input_size = POSITION_HISTORY_SIZE * size_of::<f32>();
+	let depth_size = DEPTH_SIZE * size_of::<f32>();
 	let mut inputs = Vec::with_capacity(input_size * training_samples);
+	let mut depth = Vec::with_capacity(depth_size * training_samples);
 	let mut targets = Vec::with_capacity(OUTPUT * training_samples);
 
 	// generate all frames
@@ -160,7 +162,7 @@ fn training(
 		let pos = vec3(output_event.event.position) - world_offset;
 		let angles = angles(output_event.event.angles);
 
-		g.generate_inputs(pos, angles, &mut inputs);
+		g.generate_inputs(pos, angles, &mut inputs, &mut depth);
 	}
 
 	let device = burn::backend::cuda::CudaDevice::new(gpu_id);
@@ -172,7 +174,14 @@ fn training(
 	let mut optim = AdamConfig::new().init();
 
 	let inputs = Tensor::from_data(
-		TensorData::new(inputs, Shape::new([training_samples, INPUT])),
+		TensorData::new(
+			inputs,
+			Shape::new([training_samples, POSITION_HISTORY_SIZE]),
+		),
+		&device,
+	);
+	let depth = Tensor::from_data(
+		TensorData::new(depth, Shape::new([training_samples, DEPTH_SIZE])),
 		&device,
 	);
 	let targets = Tensor::from_data(
@@ -184,7 +193,7 @@ fn training(
 	let mut best_loss = f32::INFINITY;
 
 	for epoch in 0..epochs {
-		let predictions = model.forward(inputs.clone());
+		let predictions = model.forward(inputs.clone(), depth.clone());
 
 		let loss = MseLoss::new().forward(predictions, targets.clone(), Reduction::Mean);
 
@@ -207,8 +216,7 @@ fn training(
 		model = optim.step(learning_rate, model, grads);
 
 		if epoch % (epochs >> 4) == 0 || epoch == epochs - 1 {
-			// .clone().into_scalar() extracts the f32 value from a 1-element tensor.
-			println!("  epoch {:>5} | loss = {:.8}", epoch, loss_scalar);
+			println!("  epoch {epoch:>5} | loss = {loss_scalar:.8} | best_loss = {best_loss:.8}");
 		}
 	}