Merge pull request #30 from HumanCompatibleAI/fix_latent_squeeze

Fix latent squeezing

Author: dfilan

src/reward_preprocessing/interpret.py

@@ -290,7 +290,8 @@ def param_f():
             )
             # Now, we put the latent vector thru the generator to produce transition
             # tensors that we can get observations, actions, etc out of
-            opt_latent = np.squeeze(opt_latent)
+            opt_latent = np.squeeze(opt_latent, axis=(1, 2))
+            # ^ squeeze out extraneous "height" and "width" dimensions
             opt_latent_th = th.from_numpy(opt_latent).to(th.device(device))
             opt_transitions = gan.generator(opt_latent_th)
             obs, acts, next_obs = tensor_to_transition(opt_transitions)