model init

Author: will-maclean

examples/sac_pendulum.rs

@@ -53,7 +53,7 @@ fn main() {
 
     let offline_params = OfflineAlgParams::new()
         .with_batch_size(256)
-        .with_memory_size(50000)
+        .with_memory_size(100_000)
         .with_gamma(0.99)
         .with_n_steps(100_000)
         .with_warmup_steps(10000)
@@ -75,8 +75,11 @@ fn main() {
         true,
         None,
         Some(1e-3),
-        Some(0.01),
-        Box::new(BoxSpace::from(([-1.0, -1.0, -1.0].to_vec(), [1.0, 1.0, 1.0].to_vec()))),
+        Some(0.005),
+        Box::new(BoxSpace::from((
+            vec![-1.0, -1.0, -1.0],
+            vec![1.0, 1.0, 1.0],
+        ))),
         Box::new(BoxSpace::from(([-1.0].to_vec(), [1.0].to_vec()))),
     );
 

src/common/distributions/action_distribution.rs

@@ -68,8 +68,18 @@ impl<B: Backend> DiagGaussianDistribution<B> {
         let dist: Normal<B, 2> = Normal::new(loc, std);
 
         Self {
-            means: LinearConfig::new(latent_dim, action_dim).init(device),
-            log_std: LinearConfig::new(latent_dim, action_dim).init(device),
+            means: LinearConfig::new(latent_dim, action_dim)
+                .with_initializer(burn::nn::Initializer::Uniform {
+                    min: -3e-3,
+                    max: 3e-3,
+                })
+                .init(device),
+            log_std: LinearConfig::new(latent_dim, action_dim)
+                .with_initializer(burn::nn::Initializer::Uniform {
+                    min: -3e-3,
+                    max: 3e-3,
+                })
+                .init(device),
             dist,
         }
     }

src/common/utils/module_update.rs

@@ -134,7 +134,7 @@ mod test {
         let to = Param::from_tensor(Tensor::from_floats([0.0], &Default::default()));
         let tau = 0.05;
 
-        let new_to: Param<Tensor<B, 2, Float>> = soft_update_tensor(&from, to, tau);
+        let new_to: Param<Tensor<B, 1, Float>> = soft_update_tensor(&from, to, tau);
 
         let new_to_f: f32 = new_to.val().into_scalar();
 

src/common/utils/modules.rs

@@ -21,7 +21,7 @@ impl<B: Backend> MLP<B> {
             layers.push(
                 LinearConfig::new(sizes[i], sizes[i + 1])
                     .with_initializer(burn::nn::Initializer::Uniform {
-                        min: 3e-3,
+                        min: -3e-3,
                         max: 3e-3,
                     })
                     .init(device),

src/sac/agent.rs

@@ -371,9 +371,9 @@ impl<B: AutodiffBackend> Agent<B, Vec<f32>, Vec<f32>> for SACAgent<B> {
 
         // train entropy coeficient if required to do so
         let t_ent0 = std::time::Instant::now();
-        let (ent_coef, ent_coef_loss) = self
-            .ent_coef
-            .train_step(log_prob.clone().flatten(0, 1), self.ent_lr, train_device);
+        let (ent_coef, ent_coef_loss) =
+            self.ent_coef
+                .train_step(log_prob.clone().flatten(0, 1), self.ent_lr, train_device);
         self.profiler
             .record("ent_coef", t_ent0.elapsed().as_secs_f64());
 
@@ -427,8 +427,6 @@ impl<B: AutodiffBackend> Agent<B, Vec<f32>, Vec<f32>> for SACAgent<B> {
             self.last_update = global_step;
         }
 
-        // panic!();
-
         (None, log_dict)
     }