change depth mapping

Cargo.lock

@@ -5602,9 +5602,9 @@ dependencies = [
 
 [[package]]
 name = "strafesnet_physics"
-version = "0.0.2-surf"
+version = "0.0.2-surf2"
 source = "sparse+https://git.itzana.me/api/packages/strafesnet/cargo/"
-checksum = "005a9514d77977f0bd4259c0e4c3014bfa74dbb64eaf2a00875ce37f4704880e"
+checksum = "9d86d8bd1b3af672122cd8191f18503a74674d64aa9ec26bd3e161009bda1bea"
 dependencies = [
  "arrayvec",
  "glam",

Cargo.toml

@@ -12,7 +12,7 @@ wgpu = "29.0.0"
 
 strafesnet_common = { version = "0.9.0", registry = "strafesnet" }
 strafesnet_graphics = { version = "=0.0.11-depth2", registry = "strafesnet" }
-strafesnet_physics = { version = "=0.0.2-surf", registry = "strafesnet" }
+strafesnet_physics = { version = "=0.0.2-surf2", registry = "strafesnet" }
 strafesnet_roblox_bot_file = { version = "0.9.4", registry = "strafesnet" }
 strafesnet_roblox_bot_player = { version = "=0.6.2-depth2", registry = "strafesnet" }
 strafesnet_snf = { version = "0.4.0", registry = "strafesnet" }

src/inference.rs

@@ -49,7 +49,7 @@ impl SimulateSubcommand {
 use burn::prelude::*;
 
 use crate::inputs::InputGenerator;
-use crate::net::{DEPTH_SIZE, InferenceBackend, Net, POSITION_HISTORY_SIZE};
+use crate::net::{InferenceBackend, Net, POSITION_HISTORY, POSITION_HISTORY_SIZE, SIZE};
 
 use strafesnet_common::instruction::TimedInstruction;
 use strafesnet_common::mouse::MouseState;
@@ -74,6 +74,13 @@ impl FrameState {
 			.to_array()
 			.into()
 	}
+	fn vel(&self, time: PhysicsTime) -> glam::Vec3 {
+		self.trajectory
+			.extrapolated_velocity(time)
+			.map(Into::<f32>::into)
+			.to_array()
+			.into()
+	}
 	fn angles(&self) -> glam::Vec2 {
 		self.camera.simulate_move_angles(glam::IVec2::ZERO)
 	}
@@ -178,6 +185,7 @@ fn inference(
 		let frame_state = session.get_frame_state();
 		g.generate_inputs(
 			frame_state.pos(time) - start_offset,
+			frame_state.vel(time),
 			frame_state.angles(),
 			&mut input_floats,
 			&mut depth_floats,
@@ -185,11 +193,17 @@ fn inference(
 
 		// inference
 		let inputs = Tensor::from_data(
-			TensorData::new(input_floats.clone(), Shape::new([1, POSITION_HISTORY_SIZE])),
+			TensorData::new(
+				input_floats.clone(),
+				Shape::new([1, POSITION_HISTORY_SIZE * POSITION_HISTORY]),
+			),
 			&device,
 		);
 		let depth = Tensor::from_data(
-			TensorData::new(depth_floats.clone(), Shape::new([1, DEPTH_SIZE])),
+			TensorData::new(
+				depth_floats.clone(),
+				Shape::new([1, 1, SIZE.y as usize, SIZE.x as usize]),
+			),
 			&device,
 		);
 		let outputs = model

src/inputs.rs

@@ -2,7 +2,7 @@ const LIMITS: wgpu::Limits = wgpu::Limits::defaults();
 const FORMAT: wgpu::TextureFormat = wgpu::TextureFormat::Rgba8UnormSrgb;
 use strafesnet_graphics::setup;
 
-use crate::net::{POSITION_HISTORY, SIZE};
+use crate::net::{HistoricState, POSITION_HISTORY, POSITION_HISTORY_SIZE, SIZE};
 
 // bytes_per_row needs to be a multiple of 256.
 const STRIDE_SIZE: u32 = (SIZE.x * size_of::<f32>() as u32).next_multiple_of(256);
@@ -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, glam::Vec2)>,
+	position_history: Vec<HistoricState>,
 }
 impl InputGenerator {
 	pub fn new(map: &strafesnet_common::map::CompleteMap) -> Self {
@@ -80,31 +80,35 @@ impl InputGenerator {
 	pub fn generate_inputs(
 		&mut self,
 		pos: glam::Vec3,
+		vel: 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, 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());
-			}
+		let camera =
+			strafesnet_graphics::graphics::view_inv(pos, glam::vec2(angles.x, 0.0)).inverse();
+		for state in self.position_history.iter().rev() {
+			let relative_pos = camera.project_point3(state.pos);
+			let relative_vel = camera.transform_vector3(state.vel);
+			let relative_ang = glam::vec2(angles.x - state.angles.x, state.angles.y);
+			inputs.extend_from_slice(&relative_pos.to_array());
+			inputs.extend_from_slice(&relative_vel.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, 0.0, 0.0]);
-		}
+		inputs.extend(core::iter::repeat_n(
+			0.0,
+			POSITION_HISTORY_SIZE * (POSITION_HISTORY - self.position_history.len()),
+		));
 
 		// track position history
 		if self.position_history.len() < POSITION_HISTORY {
-			self.position_history.push((pos, angles));
+			self.position_history
+				.push(HistoricState { pos, vel, angles });
 		} else {
 			self.position_history.rotate_left(1);
-			*self.position_history.last_mut().unwrap() = (pos, angles);
+			*self.position_history.last_mut().unwrap() = HistoricState { pos, vel, angles };
 		}
 
 		let mut encoder = self
@@ -165,7 +169,7 @@ impl InputGenerator {
 				self.texture_data[(STRIDE_SIZE * y) as usize
 					..(STRIDE_SIZE * y + SIZE.x * size_of::<f32>() as u32) as usize]
 					.chunks_exact(4)
-					.map(|b| 1.0 - 2.0 * f32::from_le_bytes(b.try_into().unwrap())),
+					.map(|b| 2.0 * (1.0 - f32::from_le_bytes(b.try_into().unwrap()))),
 			)
 		}
 

src/main.rs

@@ -1,3 +1,5 @@
+#![recursion_limit = "256"]
+
 use clap::{Parser, Subcommand};
 
 mod inference;

src/net.rs

@@ -1,5 +1,7 @@
 use burn::backend::Autodiff;
-use burn::nn::{Dropout, DropoutConfig, Linear, LinearConfig, Relu};
+use burn::nn::conv::{Conv2d, Conv2dConfig};
+use burn::nn::pool::{MaxPool2d, MaxPool2dConfig};
+use burn::nn::{Dropout, DropoutConfig, Linear, LinearConfig, PaddingConfig2d, Relu};
 use burn::prelude::*;
 
 pub type InferenceBackend = burn::backend::Cuda<f32>;
@@ -8,8 +10,11 @@ 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 POSITION_HISTORY_SIZE: usize = POSITION_HISTORY * 5;
-const INPUT: usize = DEPTH_SIZE + POSITION_HISTORY_SIZE;
+pub const POSITION_HISTORY_SIZE: usize = size_of::<HistoricState>() / size_of::<f32>();
+const CONV1_SIZE: usize = 8;
+const CONV2_SIZE: usize = 16;
+const INPUT: usize = ((SIZE.x >> 2) * (SIZE.y >> 2)) as usize * CONV2_SIZE
+	+ POSITION_HISTORY_SIZE * POSITION_HISTORY;
 pub const HIDDEN: [usize; 3] = [INPUT >> 3, INPUT >> 5, INPUT >> 7];
 // MoveForward
 // MoveLeft
@@ -20,9 +25,18 @@ pub const HIDDEN: [usize; 3] = [INPUT >> 3, INPUT >> 5, INPUT >> 7];
 // mouse_dy
 pub const OUTPUT: usize = 7;
 
+pub struct HistoricState {
+	pub pos: glam::Vec3,
+	pub vel: glam::Vec3,
+	pub angles: glam::Vec2,
+}
+
 #[derive(Module, Debug)]
 pub struct Net<B: Backend> {
 	input: Linear<B>,
+	conv1: Conv2d<B>,
+	conv2: Conv2d<B>,
+	pool: MaxPool2d,
 	dropout: Dropout,
 	hidden: [Linear<B>; HIDDEN.len() - 1],
 	output: Linear<B>,
@@ -40,17 +54,35 @@ impl<B: Backend> Net<B> {
 			layer
 		});
 		let output = LinearConfig::new(last_size, OUTPUT).init(device);
+		let conv1 = Conv2dConfig::new([1, CONV1_SIZE], [3, 3])
+			.with_padding(PaddingConfig2d::Same)
+			.init(device);
+		let conv2 = Conv2dConfig::new([CONV1_SIZE, CONV2_SIZE], [3, 3])
+			.with_padding(PaddingConfig2d::Same)
+			.init(device);
+		let pool = MaxPool2dConfig::new([2, 2]).with_strides([2, 2]).init();
 		let dropout = DropoutConfig::new(0.1).init();
 		Self {
 			input,
+			conv1,
+			conv2,
+			pool,
 			dropout,
 			hidden,
 			output,
 			activation: Relu::new(),
 		}
 	}
-	pub fn forward(&self, input: Tensor<B, 2>, depth: Tensor<B, 2>) -> Tensor<B, 2> {
-		let x = self.dropout.forward(depth);
+	pub fn forward(&self, input: Tensor<B, 2>, depth: Tensor<B, 4>) -> Tensor<B, 2> {
+		let x = self.conv1.forward(depth);
+		let x = self.activation.forward(x);
+		let x = self.pool.forward(x);
+		let x = self.dropout.forward(x);
+		let x = self.conv2.forward(x);
+		let x = self.activation.forward(x);
+		let x = self.pool.forward(x);
+		let x = self.dropout.forward(x);
+		let x = x.flatten(1, 3);
 		let x = Tensor::cat(vec![input, x], 1);
 		let x = self.input.forward(x);
 		let mut x = self.activation.forward(x);

src/training.rs

@@ -40,7 +40,9 @@ use burn::optim::{AdamConfig, GradientsParams, Optimizer};
 use burn::prelude::*;
 
 use crate::inputs::InputGenerator;
-use crate::net::{DEPTH_SIZE, Net, OUTPUT, POSITION_HISTORY_SIZE, TrainingBackend};
+use crate::net::{
+	DEPTH_SIZE, Net, OUTPUT, POSITION_HISTORY, POSITION_HISTORY_SIZE, SIZE, TrainingBackend,
+};
 
 use strafesnet_roblox_bot_file::v0;
 
@@ -160,9 +162,10 @@ fn training(
 		}
 
 		let pos = vec3(output_event.event.position) - world_offset;
+		let vel = vec3(output_event.event.velocity);
 		let angles = angles(output_event.event.angles);
 
-		g.generate_inputs(pos, angles, &mut inputs, &mut depth);
+		g.generate_inputs(pos, vel, angles, &mut inputs, &mut depth);
 	}
 
 	let device = burn::backend::cuda::CudaDevice::new(gpu_id);
@@ -176,12 +179,15 @@ fn training(
 	let inputs = Tensor::from_data(
 		TensorData::new(
 			inputs,
-			Shape::new([training_samples, POSITION_HISTORY_SIZE]),
+			Shape::new([training_samples, POSITION_HISTORY_SIZE * POSITION_HISTORY]),
 		),
 		&device,
 	);
 	let depth = Tensor::from_data(
-		TensorData::new(depth, Shape::new([training_samples, DEPTH_SIZE])),
+		TensorData::new(
+			depth,
+			Shape::new([training_samples, 1, SIZE.y as usize, SIZE.x as usize]),
+		),
 		&device,
 	);
 	let targets = Tensor::from_data(