From b1740764e3d1f6c9a5ac9f7f6b47755a65b966c0 Mon Sep 17 00:00:00 2001 From: ebuehrle <43623224+ebuehrle@users.noreply.github.com> Date: Thu, 28 Oct 2021 17:58:34 +0200 Subject: [PATCH] Parameterize number of discriminator updates per epoch --- src/gail/train.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/gail/train.py b/src/gail/train.py index 31edd46..ee2715c 100644 --- a/src/gail/train.py +++ b/src/gail/train.py @@ -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))