work

Author: will-maclean

src/sac/agent.rs

@@ -206,10 +206,21 @@ impl<B: AutodiffBackend> SACAgent<B> {
         let (next_action_sampled_raw, next_action_log_prob_raw) =
             self.pi.act_log_prob(next_states.clone());
 
-        let (next_action_sampled, action_scale, _) =
-            scale_actions_to_env(next_action_sampled_raw, &self.action_space, train_device);
-        let next_action_log_prob = next_action_log_prob_raw - action_scale.log().sum_dim(1);
+        let (next_action_sampled, action_scale, _) = scale_actions_to_env(
+            next_action_sampled_raw.clone(),
+            &self.action_space,
+            train_device,
+        );
+        // let next_action_log_prob =
+        //     next_action_log_prob_raw.clone() - action_scale.clone().log().sum_dim(1);
+
+        let next_action_log_prob_scaled = next_action_log_prob_raw
+            - (next_action_sampled_raw.powi_scalar(2).neg().add_scalar(1.0))
+                .mul(action_scale)
+                .add_scalar(1e-6)
+                .log();
 
+        let next_action_log_prob_scaled = next_action_log_prob_scaled.sum_dim(1);
         // disp_tensorf("next_action_sampled", &next_action_sampled);
         // disp_tensorf("next_action_log_prob", &next_action_log_prob);
 
@@ -225,7 +236,7 @@ impl<B: AutodiffBackend> SACAgent<B> {
         // disp_tensorf("2next_q_vals", &next_q_vals);
 
         // add the entropy term
-        let next_q_vals = next_q_vals - next_action_log_prob.mul_scalar(ent_coef);
+        let next_q_vals = next_q_vals - next_action_log_prob_scaled.mul_scalar(ent_coef);
         // disp_tensorf("3next_q_vals", &next_q_vals);
 
         // td error + entropy term
@@ -358,12 +369,16 @@ impl<B: AutodiffBackend> Agent<B, Vec<f32>, Vec<f32>> for SACAgent<B> {
         let t_policy0 = std::time::Instant::now();
         let (actions_pi_raw, log_prob_raw) = self.pi.act_log_prob(states.clone());
 
-        let (actions_pi_scaled, action_scale, action_bias) =
-            scale_actions_to_env(actions_pi_raw, &self.action_space, train_device);
+        let (actions_pi_scaled, action_scale, _) =
+            scale_actions_to_env(actions_pi_raw.clone(), &self.action_space, train_device);
 
         // let log_prob = log_prob_raw - action_scale.log().sum_dim(1);
 
-        let log_prob_scaled = log_prob - ((1 - actions_pi_raw.powi(2)) * action_scale + 1e-6).log();
+        let log_prob_scaled = log_prob_raw
+            - (actions_pi_raw.powi_scalar(2).neg().add_scalar(1.0))
+                .mul(action_scale)
+                .add_scalar(1e-6)
+                .log();
         let log_prob_scaled = log_prob_scaled.sum_dim(1);
 
         self.profiler

src/sac/models.rs

@@ -1,16 +1,13 @@
 use burn::{
     module::Module,
+    nn::{Linear, LinearConfig},
     prelude::Backend,
     tensor::{activation::relu, Tensor},
 };
 
 use crate::common::{
     agent::Policy,
-    distributions::{
-        action_distribution::{ActionDistribution, SquashedDiagGaussianDistribution},
-        distribution::BaseDistribution,
-        normal::Normal,
-    },
+    distributions::{distribution::BaseDistribution, normal::Normal},
     utils::modules::MLP,
 };
 
@@ -27,8 +24,8 @@ impl<B: Backend> PiModel<B> {
     pub fn new(obs_size: usize, n_actions: usize, device: &B::Device) -> Self {
         Self {
             mlp: MLP::new(&[obs_size, 256, 256].to_vec(), device),
-            scale_head: LinearConfig::new(256, n_actions),
-            loc_head: LinearConfig::new(256, n_actions),
+            scale_head: LinearConfig::new(256, n_actions).init(device),
+            loc_head: LinearConfig::new(256, n_actions).init(device),
             // dist: SquashedDiagGaussianDistribution::new(256, n_actions, device, 1e-6),
             n_actions,
         }
@@ -37,41 +34,42 @@ impl<B: Backend> PiModel<B> {
 
 impl<B: Backend> PiModel<B> {
     fn forward(&self, obs: Tensor<B, 2>) -> (Tensor<B, 2>, Tensor<B, 2>) {
-        let latent = relu(self.mlp.forward(obs.clone().unsqueeze_dim(0)));
+        let latent = relu(self.mlp.forward(obs.clone()));
         let loc = self.loc_head.forward(latent.clone());
         let log_scale = self.loc_head.forward(latent);
         let log_scale = log_scale.tanh();
 
-        let min_log_scale = -20;
-        let max_log_scale = 2;
+        let min_log_scale = -20.0;
+        let max_log_scale = 2.0;
 
         let log_scale = min_log_scale + 0.5 * (max_log_scale - min_log_scale) * (log_scale + 1.0);
 
         (loc, log_scale)
     }
     pub fn act(&mut self, obs: &Tensor<B, 1>, deterministic: bool) -> Tensor<B, 1> {
-        let (loc, log_scale) = self.forward(obs.unsqueeze_dim(0));
+        let (loc, log_scale) = self.forward(obs.clone().unsqueeze_dim(0));
 
         if deterministic {
-            loc.tanh()
+            loc.tanh().squeeze(0)
         } else {
             let scale = log_scale.exp();
             let dist = Normal::new(loc, scale);
             let x_t = dist.rsample();
             let action = x_t.tanh();
 
-            action
+            action.squeeze(0)
         }
 
         // self.dist.actions_from_obs(latent, deterministic).squeeze(0)
     }
 
     pub fn act_log_prob(&mut self, obs: Tensor<B, 2>) -> (Tensor<B, 2>, Tensor<B, 2>) {
         let (loc, log_scale) = self.forward(obs.unsqueeze_dim(0));
+        let scale = log_scale.exp();
         let dist = Normal::new(loc, scale);
         let x_t = dist.rsample();
-        let action = x_t.tanh();
-        let log_prob = dist.log_prob(action);
+        let action = x_t.clone().tanh();
+        let log_prob = dist.log_prob(x_t);
 
         (action, log_prob)