Merge pull request #26 from HumanCompatibleAI/harmonize_naming

Harmonize reward net naming conventions

Author: PavelCz

src/reward_preprocessing/common/utils.py

@@ -221,17 +221,17 @@ class RewardGeneratorCombo(nn.Module):
     Assumes that the RewardNet normalizes observations to [0,1].
     """
 
-    def __init__(self, reward_net: RewardNet, generator: nn.Module):
+    def __init__(self, rew_net: RewardNet, generator: nn.Module):
         super().__init__()
-        self.reward_net = reward_net
+        self.rew_net = rew_net
         self.generator = generator
 
     def forward(self, latent_tens: th.Tensor):
         latent_vec = th.mean(latent_tens, dim=[2, 3])
         transition_tensor = self.generator(latent_vec)
         obs, action_vec, next_obs = tensor_to_transition(transition_tensor)
         done = th.zeros(action_vec.shape)
-        return self.reward_net.forward(obs, action_vec, next_obs, done)
+        return self.rew_net.forward(obs, action_vec, next_obs, done)
 
 
 def log_img_wandb(

src/reward_preprocessing/interpret.py

@@ -179,7 +179,7 @@ def interpret(
         # Combine rew net with GAN.
         gan = th.load(gan_path, map_location=th.device(device))
         model_to_analyse = RewardGeneratorCombo(
-            reward_net=rew_net, generator=gan.generator
+            rew_net=rew_net, generator=gan.generator
         )
 
     model_to_analyse.eval()  # Eval for visualization.