Adds support for log-std parameter in ActorCritic (#67)

* Fixes gradient propogation through std-dev

* adds support for log std

* renames to noise_std_type

* fixes value error

* adds back option for noise std type

Author: Mayankm96

config/dummy_config.yaml

@@ -71,6 +71,7 @@ policy:
   actor_hidden_dims: [128, 128, 128]
   critic_hidden_dims: [128, 128, 128]
   init_noise_std: 1.0
+  noise_std_type: "scalar"  # 'scalar' or 'log'
   # only needed for `ActorCriticRecurrent`
   # rnn_type: 'lstm'
   # rnn_hidden_size: 512

rsl_rl/modules/actor_critic.py

@@ -24,6 +24,7 @@ def __init__(
         critic_hidden_dims=[256, 256, 256],
         activation="elu",
         init_noise_std=1.0,
+        noise_std_type: str = "scalar",
         **kwargs,
     ):
         if kwargs:
@@ -64,7 +65,15 @@ def __init__(
         print(f"Critic MLP: {self.critic}")
 
         # Action noise
-        self.std = nn.Parameter(init_noise_std * torch.ones(num_actions))
+        self.noise_std_type = noise_std_type
+        if self.noise_std_type == "scalar":
+            self.std = nn.Parameter(init_noise_std * torch.ones(num_actions))
+        elif self.noise_std_type == "log":
+            self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions)))
+        else:
+            raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'")
+
+        # Action distribution (populated in update_distribution)
         self.distribution = None
         # disable args validation for speedup
         Normal.set_default_validate_args(False)
@@ -100,8 +109,16 @@ def entropy(self):
         return self.distribution.entropy().sum(dim=-1)
 
     def update_distribution(self, observations):
+        # compute mean
         mean = self.actor(observations)
-        std = self.std.expand_as(mean)
+        # compute standard deviation
+        if self.noise_std_type == "scalar":
+            std = self.std.expand_as(mean)
+        elif self.noise_std_type == "log":
+            std = torch.exp(self.log_std).expand_as(mean)
+        else:
+            raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'")
+        # create distribution
         self.distribution = Normal(mean, std)
 
     def act(self, observations, **kwargs):

rsl_rl/runners/on_policy_runner.py

@@ -261,7 +261,7 @@ def log(self, locs: dict, width: int = 80, pad: int = 35):
                 else:
                     self.writer.add_scalar("Episode/" + key, value, locs["it"])
                     ep_string += f"""{f'Mean episode {key}:':>{pad}} {value:.4f}\n"""
-        mean_std = self.alg.actor_critic.std.mean()
+        mean_std = self.alg.actor_critic.action_std.mean()
         fps = int(self.num_steps_per_env * self.env.num_envs / (locs["collection_time"] + locs["learn_time"]))
 
         # -- Losses