huangfu/InteractionImitation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge branch 'main' into dev

Browse Source
This commit is contained in:
ebuehrle
2021-10-28 18:01:02 +02:00
parent 8d7409c914 b1740764e3
commit 5bdf85562b
12 changed files with 353 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/gail/options.py
View File

@@ -40,7 +40,7 @@ class OptionsEnv(gym.Wrapper):
# action 0 is considered safe fallback
self.m[0] = True
self.ch, self.value, self.log_prob = generator.policy.predict({
self.ch, self.value, self.log_prob = generator.policy.forward({
'obs': torch.tensor(self.s).unsqueeze(0).to(generator.policy.device),
'mask': self.m.unsqueeze(0).to(generator.policy.device),
})

5
src/gail/train.py
View File

@@ -9,10 +9,11 @@ def flatten_transitions(transitions):
'dones': np.stack(list(t['dones'] for t in transitions), axis=0),
}
def train_discriminator(env, generator, discriminator, num_samples):
def train_discriminator(env, generator, discriminator, num_samples, n_updates=1):
transitions = list(itertools.islice(env.sample_ll(generator), num_samples))
generator_samples = flatten_transitions(transitions)
discriminator.train_disc(gen_samples=generator_samples)
for _ in range(n_updates):
discriminator.train_disc(gen_samples=generator_samples)
def train_generator(env, generator, discriminator, num_samples):
generator_samples = list(itertools.islice(env.sample_hl(generator, discriminator), num_samples+1))