Implement ValueDICE and some restructuring

This commit is contained in:
Johannes Fischer
2021-08-04 19:34:49 +02:00
parent cba42c6e4d
commit 5c40de66fa
8 changed files with 442 additions and 86 deletions

9
src/util/nn_training.py Normal file
View File

@@ -0,0 +1,9 @@
def optimizer_factory(config, parameters):
optimizer_type = config['optimizer']
learning_rate = config['lr']
weight_decay = config['weight_decay']
if optimizer_type == 'adam':
optimizer = torch.optim.Adam(parameters, lr=learning_rate, weight_decay=weight_decay)
else:
raise NotImplementedError
return optimizer