Fix discriminator reward
Had wrong sign.
This commit is contained in:
@@ -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)
|
||||||
|
|||||||
@@ -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
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user