Add readme.md

readme.md

@@ -0,0 +1 @@
+**請助教看`tank-rust/dqn/`下的程式**

tank-rust/src/dqn/collect.rs

@@ -10,14 +10,14 @@ use burn::{backend::wgpu::WgpuDevice, module::Module, record::NoStdTrainingRecor
 use rand::{thread_rng, Rng};
 
 use super::model::{DQNModel, DQNModelConfig};
-const EXPLORE_RATE: f32 = 0.8;
-
 pub struct App<'a> {
     model: DQNModel<Backend>,
     device: WgpuDevice,
     last_state_action: Option<(Info<'a>, Action)>,
     #[cfg(feature = "train")]
     outlet: BufWriter<File>,
+    #[cfg(feature = "train")]
+    explore_rate: f32,
 }
 
 impl<'a> App<'a> {
@@ -33,6 +33,17 @@ impl<'a> App<'a> {
                 &device,
             )
             .unwrap();
+        #[cfg(feature = "train")]
+        let explore_rate = std::env::var("EPSILON")
+            .map(|x| {
+                let n: usize = x.parse().ok()?;
+                Some(1.0 / (n as f32 + 2.0).log2() - 0.03)
+            })
+            .into_iter()
+            .flatten()
+            .next()
+            .unwrap_or(0.4);
+
         Self {
             model,
             device,
@@ -45,6 +56,8 @@ impl<'a> App<'a> {
                     .open(format!("{model_path}/dataset"))
                     .unwrap(),
             ),
+            #[cfg(feature = "train")]
+            explore_rate,
         }
     }
     #[cfg(feature = "train")]
@@ -64,8 +77,8 @@ impl<'a> App<'a> {
             bincode::serialize_into(&mut self.outlet, &item).unwrap();
         }
 
-        let action = match thread_rng().gen_ratio((4096.0 * EXPLORE_RATE) as u32, 4096) {
-            true => match thread_rng().gen_range(0..(ACTION_SPACE+2) as i32) {
+        let action = match thread_rng().gen_ratio((4096.0 * self.explore_rate) as u32, 4096) {
+            true => match thread_rng().gen_range(0..(ACTION_SPACE + 2) as i32) {
                 0 => Action::TurnRight,
                 1 => Action::TurnLeft,
                 2 => Action::AimRight,
@@ -83,7 +96,7 @@ impl<'a> App<'a> {
     pub fn predict_action(&self, state: &Info) -> Action {
         let input = state.into_feature_tensor(&self.device).unsqueeze(); // Convert input tensor to shape [1, input_size]
         let ans = self.model.forward(input);
-        match ans.argmax(1).into_scalar(){
+        match ans.argmax(1).into_scalar() {
             0 => Action::TurnRight,
             1 => Action::TurnLeft,
             2 => Action::AimRight,

tank-rust/src/dqn/feature.rs

@@ -18,7 +18,7 @@ impl Polar {
     pub fn clip(&self) -> Self {
         Polar {
             angle: self.angle,
-            distance: self.distance.min(1e6).max(0.0),
+            distance: self.distance.max(0.0).min(1e3),
         }
     }
 }
@@ -118,9 +118,6 @@ impl<'a> Info<'a> {
         let wall = self
             .player
             .closest(self.walls.iter().map(|wall| (wall.x, wall.y)));
-        let bullet = self
-            .player
-            .closest(self.bullets.iter().map(|bullet| (bullet.x, bullet.y)));
 
         let target = self.get_target().get_pos(self).clip();
 
@@ -131,7 +128,7 @@ impl<'a> Info<'a> {
             normalize_angle(target.angle - angle).tanh(),
             (wall.distance - target.distance).tanh(),
             (self.player.power as f32).tanh(),
-            (wall.distance + 1.0).log2(),
+            (wall.clip().distance + 1.0).log2(),
             (emeny.distance + 1.0).log2(),
             normalize_angle(emeny.angle - gun_angle).tanh(),
             normalize_angle(wall.angle - gun_angle).tanh(),
@@ -139,6 +136,7 @@ impl<'a> Info<'a> {
     }
     pub fn into_feature_tensor<B: Backend>(&self, device: &B::Device) -> Tensor<B, 1> {
         let feature = self.into_feature();
+
         Tensor::from_floats(feature, device)
     }
     fn get_target(&self) -> Target {
@@ -179,7 +177,7 @@ impl<'a> Info<'a> {
 
         reward
             + match next.player.score - self.player.score {
-                x if x > 2 => 20.0,
+                x if x > 2 => 20.0, // bypass emeny
                 x if x > 0 => 10.0, // too high, tank my ignore power station
                 _ => -1.0,
             }

tank-rust/src/dqn/mod.rs

@@ -6,7 +6,7 @@ mod training;
 
 pub mod prelude {
     pub use super::collect::App as DQNApp;
-    pub use super::dataset::{TankDataset, TankItem};
+    pub use super::dataset::{TankBatcher, TankDataset, TankItem};
     pub use super::feature::{ACTION_SPACE, FEATRUE_SPACE};
-    pub use super::training::run as train;
+    pub use super::training::{run as train, ExpConfig};
 }

tank-rust/src/fit.rs

@@ -1,77 +1,95 @@
-use burn::data::dataset::Dataset;
+use burn::{
+    data::dataloader::DataLoaderBuilder,
+    optim::{AdamConfig, SgdConfig},
+    record::{CompactRecorder, NoStdTrainingRecorder},
+    tensor::backend::AutodiffBackend,
+    train::{
+        metric::{
+            store::{Aggregate, Direction, Split},
+            LossMetric,
+        },
+        LearnerBuilder, MetricEarlyStoppingStrategy, StoppingCondition,
+    },
+};
 
-use crate::dqn::prelude::TankItem;
-use crate::ffi::prelude::*;
-use rand::Rng;
+use crate::dqn::prelude::{ExpConfig, TankBatcher, TankDataset};
 
-// fn random_action() -> Action {
-//     let mut rng = rand::thread_rng();
-//     match rng.gen_range(0..2) {
-//         0 => Action::AimLeft,
-//         1 => Action::Forward,
-//         _ => unreachable!(),
-//     }
-// }
+pub fn run<B: AutodiffBackend>(device: B::Device) {
+    // let d = [
+    //     feature[0],
+    //     -feature[0],
+    //     shoot_target_angle*0.7*feature[2],
+    //     -shoot_target_angle*0.7*feature[2],
+    //     8.0 * feature[2] / shoot_target_distance / shoot_target_angle,
+    //     feature[2]*shoot_target_distance*0.3-feature[2],
+    // ];
 
-// fn random_item() -> TankItem {
-//     let mut previous_info=Info::default();
-//     TankItem {
-//         previous_state: todo!(),
-//         new_state: todo!(),
-//         action: todo!(),
-//         reward: todo!(),
-//     }
-// }
+    let optimizer = AdamConfig::new();
+    let config = ExpConfig::new(optimizer);
+    let mut model = DQNModelConfig::new().init(&device);
 
-pub struct FitDataset;
-
-impl FitDataset {
-    /// Get closer to the power station
-    fn close_power_station() -> TankItem {
-        let mut power_stations = Station::default();
-
-        let mut previous_info = Info::default();
-        let mut new_info = Info::default();
-        let mut rng = rand::thread_rng();
-        previous_info.player.power = rng.gen_range(0..2);
-        new_info.player.power = previous_info.player.power;
-        previous_info.player.angle = rng.gen_range(0..360);
-        new_info.player.angle = previous_info.player.angle;
-
-        TankItem {
-            previous_state: todo!(),
-            new_state: todo!(),
-            action: Action::Forward,
-            reward: todo!(),
-        }
-    }
-    /// Flee from power station if power is high
-    fn flee_power_station() -> TankItem {
-        let mut previous_info = Info::default();
-        TankItem {
-            previous_state: todo!(),
-            new_state: todo!(),
-            action: Action::Backward,
-            reward: todo!(),
-        }
-    }
-}
-
-impl Dataset<TankItem> for FitDataset {
-    fn get(&self, _: usize) -> Option<TankItem> {
-        let previous_state = todo!();
-        let new_state = todo!();
-        let action = Action::AimLeft;
-        let reward = 0.0;
-        Some(TankItem {
-            previous_state,
-            new_state,
-            action,
-            reward,
-        })
-    }
-
-    fn len(&self) -> usize {
-        1
+    if fs::metadata(format!("{model_path}/model")).is_ok() {
+        model = model
+            .load_file(
+                format!("{model_path}/model"),
+                &NoStdTrainingRecorder::new(),
+                &device,
+            )
+            .unwrap();
     }
+
+    // Define train/test datasets and dataloaders
+
+    let train_dataset = TankDataset::train();
+    let test_dataset = TankDataset::test();
+
+    println!("Train Dataset Size: {}", train_dataset.len());
+    println!("Test Dataset Size: {}", test_dataset.len());
+
+    let batcher_train = TankBatcher::<B>::new(device.clone());
+
+    let batcher_test = TankBatcher::<B::InnerBackend>::new(device.clone());
+
+    // Since dataset size is small, we do full batch gradient descent and set batch size equivalent to size of dataset
+
+    let dataloader_train = DataLoaderBuilder::new(batcher_train)
+        .batch_size(train_dataset.len())
+        .shuffle(config.seed)
+        .num_workers(config.num_workers)
+        .build(train_dataset);
+
+    let dataloader_test = DataLoaderBuilder::new(batcher_test)
+        .batch_size(test_dataset.len())
+        .shuffle(config.seed)
+        .num_workers(config.num_workers)
+        .build(test_dataset);
+
+    // Model
+    let learner = LearnerBuilder::new(ARTIFACT_DIR)
+        .metric_train_numeric(LossMetric::new())
+        .metric_valid_numeric(LossMetric::new())
+        .with_file_checkpointer(CompactRecorder::new())
+        .early_stopping(MetricEarlyStoppingStrategy::new::<LossMetric<B>>(
+            Aggregate::Mean,
+            Direction::Lowest,
+            Split::Valid,
+            StoppingCondition::NoImprovementSince { n_epochs: 1 },
+        ))
+        .devices(vec![device.clone()])
+        .num_epochs(config.num_epochs)
+        .summary()
+        .build(model, config.optimizer.init(), 5e-3);
+
+    let model_trained = learner.fit(dataloader_train, dataloader_test);
+
+    config
+        .save(format!("{ARTIFACT_DIR}/config.json").as_str())
+        .unwrap();
+
+    model_trained
+        .save_file(
+            format!("{ARTIFACT_DIR}/model"),
+            &NoStdTrainingRecorder::new(),
+        )
+        .expect("Failed to save trained model");
 }

tank-rust/src/lib.rs

@@ -1,6 +1,6 @@
 mod dqn;
 mod ffi;
-use std::{ffi::OsStr, os::unix::ffi::OsStrExt};
+use std::ffi::OsString;
 
 use burn::backend::{wgpu::AutoGraphicsApi, Wgpu};
 use dqn::prelude::*;
@@ -12,11 +12,9 @@ type Backend = Wgpu<AutoGraphicsApi, f32, i32>;
 
 #[no_mangle]
 pub extern "C" fn init(model_path: *const u8, len: i32) -> *mut DQNApp<'static> {
-    let model_path =
-        OsStr::from_bytes(unsafe { std::slice::from_raw_parts(model_path, len as usize) })
-            .to_str()
-            .unwrap();
-    let app = DQNApp::new(model_path);
+    let model_path = unsafe { std::slice::from_raw_parts(model_path, len as usize) };
+    let model_path = unsafe { OsString::from_encoded_bytes_unchecked(model_path.to_vec()) };
+    let app = DQNApp::new(model_path.to_str().unwrap());
 
     Box::into_raw(Box::new(app))
 }

train.sh

@@ -9,7 +9,8 @@ for i in  $(seq 1 $1);
 do
    echo "epoch $i"
    cd TankMan
-   timeout 240 python -m mlgame -f 3000 \
+   export EPSILON=$i
+   timeout 200 python -m mlgame -f 3000 \
    -i ../ml/collect1.py -i ../ml/collect2.py \
    . --green_team_num 1 --blue_team_num 1 --is_manual 1 \
    --frame_limit 1000 > /dev/null