one bug at a time

will-maclean · Sep 1, 2024 · d924594 · d924594
1 parent 5c9ea16
commit d924594
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/src/common/distributions/action_distribution.rs b/src/common/distributions/action_distribution.rs
@@ -80,7 +80,7 @@ impl<B: Backend> DiagGaussianDistribution<B> {
             log_std: Param::from_tensor(
                 Tensor::ones(Shape::new([action_dim]), device).mul_scalar(log_std_init),
             ),
-            dist: dist.no_grad(),
+            dist,
         }
     }
 }
@@ -116,10 +116,12 @@ impl<B: Backend> ActionDistribution<B> for DiagGaussianDistribution<B> {
 
         let loc = self.means.forward(obs);
 
+        self.dist = Normal::new(loc.clone(), scale);
+
         if deterministic {
             loc
         } else {
-            Normal::new(loc, scale).sample()
+            self.dist.sample()
         }
     }
 }

diff --git a/src/env/classic_control/pendulum.rs b/src/env/classic_control/pendulum.rs
@@ -142,9 +142,9 @@ impl Env<Vec<f32>, Vec<f32>> for PendulumEnv {
 
 pub fn make_pendulum(max_steps: Option<usize>) -> Box<dyn Env<Vec<f32>, Vec<f32>>> {
     let env = make_pendulum_eval(max_steps);
-    let env = ScaleRewardWrapper::new(env, 0.01);
+    // let env = ScaleRewardWrapper::new(env, 0.01);
 
-    Box::new(env)
+    env
 }
 
 pub fn make_pendulum_eval(max_steps: Option<usize>) -> Box<dyn Env<Vec<f32>, Vec<f32>>> {