tau check, module init fn

Author: will-maclean

src/common/utils/module_update.rs

@@ -61,12 +61,12 @@ pub fn update_conv2d<B: Backend>(from: &Conv2d<B>, to: Conv2d<B>, tau: Option<f3
 mod test {
     use burn::{
         backend::NdArray,
-        module::Module,
+        module::{Module, Param},
         nn::{Linear, LinearConfig},
-        tensor::backend::Backend,
+        tensor::{backend::Backend, Float, Tensor},
     };
 
-    use crate::common::agent::Policy;
+    use crate::common::{agent::Policy, utils::module_update::soft_update_tensor};
 
     use super::update_linear;
 
@@ -126,4 +126,20 @@ mod test {
 
         a.update(&b, None);
     }
+
+    #[test]
+    fn test_soft_update_value() {
+        type B = NdArray;
+        let from = Param::from_tensor(Tensor::from_floats([1.0], &Default::default()));
+        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_f: f32 = new_to.val().into_scalar();
+
+        let diff = tau - new_to_f;
+
+        assert!(diff < 1e-6);
+    }
 }

src/common/utils/modules.rs

@@ -18,7 +18,14 @@ impl<B: Backend> MLP<B> {
         let mut layers = Vec::new();
 
         for i in 0..sizes.len() - 1 {
-            layers.push(LinearConfig::new(sizes[i], sizes[i + 1]).init(device))
+            layers.push(
+                LinearConfig::new(sizes[i], sizes[i + 1])
+                    .with_initializer(burn::nn::Initializer::Uniform {
+                        min: 3e-3,
+                        max: 3e-3,
+                    })
+                    .init(device),
+            )
         }
 
         Self { layers }