Pad and Cut Trajectory to correct length

priorConditionedAnnealing/pca.py

@@ -6,6 +6,7 @@ from torch import nn
 from stable_baselines3.common.distributions import Distribution as SB3_Distribution
 from stable_baselines3.common.distributions import sum_independent_dims
 from torch.distributions import Normal
+import torch.nn.functional as F
 
 
 class Par_Strength(Enum):
@@ -112,7 +113,7 @@ class PCA_Distribution(SB3_Distribution):
         return sum_independent_dims(self.distribution.entropy())
 
     def sample(self, traj: th.Tensor) -> th.Tensor:
-        pi_mean, pi_std = self.distribution.mean, self.distribution.scale,
+        pi_mean, pi_std = self.distribution.mean, self.distribution.scale
         rho_mean, rho_std = self._conditioning_engine(traj, pi_mean, pi_std)
         eta = self._get_rigged(pi_mean, pi_std,
                                rho_mean, rho_std)
@@ -136,7 +137,15 @@ class PCA_Distribution(SB3_Distribution):
 
         return eta.detach()
 
-    def _conditioning_engine(self, traj, pi_mean, pi_std):
+    def _pad_and_cut_trajectory(self, traj, value=0):
+        cut = traj[:self.window]
+        if traj.shape[-2] < self.window:
+            missing = self.window - traj.shape[-2]
+            return F.pad(input=cut, pad=(missing, 0), value=value)
+        return cut
+
+    def _conditioning_engine(self, trajectory, pi_mean, pi_std):
+        traj = self._pad_and_cut_trajectory(trajectory)
         y_np = np.append(np.swapaxes(traj, -1, -2),
                          np.expand_dims(pi_mean, -1), -1)