Experimental Support for PCA

2023-05-21 14:27:09 +02:00
parent 76ea3a6326
commit 2efa6c18fb
2 changed files with 25 additions and 3 deletions
@@ -213,7 +213,10 @@ class GaussianRolloutCollectorAuxclass():
            with th.no_grad():
                # Convert to pytorch tensor or to TensorDict
                obs_tensor = obs_as_tensor(self._last_obs, self.device)
-                actions, values, log_probs = self.policy(obs_tensor)
+                if 'use_pca' in self.policy and self.policy['use_pca']:
+                    actions, values, log_probs = self.policy(obs_tensor, trajectory=self.get_past_trajectories())
+                else:
+                    actions, values, log_probs = self.policy(obs_tensor)
                dist = self.policy.get_distribution(obs_tensor).distribution
                mean, chol = get_mean_and_chol(dist)
            actions = actions.cpu().numpy()
@@ -274,3 +277,7 @@ class GaussianRolloutCollectorAuxclass():
        callback.on_rollout_end()

        return True
+
+    def get_past_trajectories(self):
+        # TODO: Respect Episode Boundaries
+        return self.actions