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

Merge branch 'main' of github.com:sisl/InteractionImitation

Browse Source
This commit is contained in:
Johannes Fischer
2021-10-28 18:17:47 +02:00
parent e36aaef9f2 b1740764e3
commit 1a1f6d8836
1 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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))