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

Parameterize number of discriminator updates per epoch

Browse Source
This commit is contained in:
ebuehrle
2021-10-28 17:58:34 +02:00
parent 070b8fc785
commit b1740764e3
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), '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)) transitions = list(itertools.islice(env.sample_ll(generator), num_samples))
generator_samples = flatten_transitions(transitions) 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): def train_generator(env, generator, discriminator, num_samples):
generator_samples = list(itertools.islice(env.sample_hl(generator, discriminator), num_samples+1)) generator_samples = list(itertools.islice(env.sample_hl(generator, discriminator), num_samples+1))