Fix: make act_inference return policy mean (without std dev) at deployment time (#118)

* fix to make act_inference return just the policy mean at deployment time, when the policy was trained with state-dependent standard-dev.

Author: iakinola23

rsl_rl/modules/actor_critic.py

@@ -148,7 +148,10 @@ def act(self, obs, **kwargs):
     def act_inference(self, obs):
         obs = self.get_actor_obs(obs)
         obs = self.actor_obs_normalizer(obs)
-        return self.actor(obs)
+        if self.state_dependent_std:
+            return self.actor(obs)[..., 0, :]
+        else:
+            return self.actor(obs)
 
     def evaluate(self, obs, **kwargs):
         obs = self.get_critic_obs(obs)

rsl_rl/modules/actor_critic_recurrent.py

@@ -167,7 +167,10 @@ def act_inference(self, obs):
         obs = self.get_actor_obs(obs)
         obs = self.actor_obs_normalizer(obs)
         out_mem = self.memory_a(obs).squeeze(0)
-        return self.actor(out_mem)
+        if self.state_dependent_std:
+            return self.actor(out_mem)[..., 0, :]
+        else:
+            return self.actor(out_mem)
 
     def evaluate(self, obs, masks=None, hidden_states=None):
         obs = self.get_critic_obs(obs)