Fix discriminator reward

Had wrong sign.
This commit is contained in:
ebuehrle
2021-09-15 07:15:43 +02:00
parent 183657dc36
commit 4928458e08
2 changed files with 9 additions and 5 deletions

View File

@@ -36,7 +36,8 @@ class CnnDiscriminator(torch.nn.Module):
def forward(self, state, action): def forward(self, state, action):
sa = self._concatenate(state, action) sa = self._concatenate(state, action)
return self.cnn(sa).squeeze() assert sa.ndim == 4
return self.cnn(sa).squeeze(1)
class MlpDiscriminator(torch.nn.Module): class MlpDiscriminator(torch.nn.Module):
"""MLP similar to stable_baselines3.common.policies.ActorCriticPolicy.""" """MLP similar to stable_baselines3.common.policies.ActorCriticPolicy."""
@@ -55,4 +56,5 @@ class MlpDiscriminator(torch.nn.Module):
def forward(self, state, action): def forward(self, state, action):
flat = self.flatten(state) flat = self.flatten(state)
sa = torch.cat((action, flat), -1) sa = torch.cat((action, flat), -1)
return self.mlp(sa).squeeze() assert sa.ndim == 2
return self.mlp(sa).squeeze(1)

View File

@@ -135,9 +135,11 @@ class HLOptions(OptionsEnv):
self.steps = 0 self.steps = 0
def _after_step(self): def _after_step(self):
self.r += self.discount**self.steps * self.discriminator.discrim_net.discriminator( self.r += self.discount**self.steps * self.discriminator.discrim_net.reward_train(
torch.tensor(self.s).unsqueeze(0).to(self.discriminator.discrim_net.device()), state=torch.tensor(self.s).unsqueeze(0).to(self.discriminator.discrim_net.device()),
torch.tensor([[self.a]]).to(self.discriminator.discrim_net.device()), action=torch.tensor([[self.a]]).to(self.discriminator.discrim_net.device()),
next_state=torch.tensor(self.s).unsqueeze(0).to(self.discriminator.discrim_net.device()), # unused
done=torch.tensor(self.done).unsqueeze(0).to(self.discriminator.discrim_net.device()), # unused
) )
self.steps += 1 self.steps += 1