From bf4c19a4d0869b7702a260889b5338c2f217bcb4 Mon Sep 17 00:00:00 2001 From: Johannes Fischer Date: Thu, 5 Aug 2021 18:33:34 +0200 Subject: [PATCH] Add value dice ray config --- experiments/experiment.py | 22 ++++++++++ src/value_dice/value_dice.py | 83 +++++++++++++++++++++++++----------- 2 files changed, 80 insertions(+), 25 deletions(-) diff --git a/experiments/experiment.py b/experiments/experiment.py index 40235d0..bcddb06 100644 --- a/experiments/experiment.py +++ b/experiments/experiment.py @@ -90,6 +90,9 @@ def get_full_config(ray_config:dict, method:str)->dict: if method == 'bc': from src.bc import bc_config config = bc_config(ray_config) + elif method == 'vd': + from src.value_dice import vd_config + config = vd_config(ray_config) else: raise NotImplementedError return config @@ -118,6 +121,25 @@ def get_ray_config(method:str)->dict: "head_hidden_dim": tune.lograndint(16,257), "head_final_activation": tune.choice(['sigmoid', None]), } + elif method == 'vd': + ray_config = { + "policy_lr": tune.loguniform(1e-5, 1e-3), + "value_lr": tune.loguniform(1e-5, 1e-3), + "policy_weight_decay": tune.choice([0, 0.1]), + "value_weight_decay": tune.choice([0, 0.1]), + "train_batch_size": tune.choice([16,32,64]), + "deepsets_phi_hidden_n": tune.randint(1,5), + "deepsets_phi_hidden_dim": tune.lograndint(8,65), + "deepsets_latent_dim": tune.lograndint(8,129), + "deepsets_rho_hidden_n": tune.randint(0,3), + "deepsets_rho_hidden_dim": tune.lograndint(8,129), + "deepsets_output_dim": tune.lograndint(4,129), + "head_hidden_n": tune.randint(1,6), + "head_hidden_dim": tune.lograndint(16,257), + "head_final_activation": tune.choice(['sigmoid', None]), + "clip_grad_norm": tune.choice([.5, 1., 5., 10.]), + "discount": tune.choice([.95, .99]) + } else: raise NotImplementedError return ray_config diff --git a/src/value_dice/value_dice.py b/src/value_dice/value_dice.py index 682f40e..6f9ee16 100644 --- a/src/value_dice/value_dice.py +++ b/src/value_dice/value_dice.py @@ -11,39 +11,72 @@ from src.util.nn_training import optimizer_factory from tqdm import tqdm import json5 from ray import tune + def vd_config(ray_config): config = { - 'ego_encoder': {'input_dim': 5, 'hidden_n': 0, 'hidden_dim':0, 'output_dim': 0}, - 'deepsets': { - 'input_dim': 6, - 'phi': { - 'hidden_n': ray_config['deepsets_phi_hidden_n'], - 'hidden_dim': ray_config['deepsets_phi_hidden_dim'] - }, - 'latent_dim': ray_config['deepsets_latent_dim'], - 'rho': { - 'hidden_n': ray_config['deepsets_rho_hidden_n'], - 'hidden_dim': ray_config['deepsets_rho_hidden_dim'] - }, - 'output_dim': ray_config['deepsets_output_dim'] + 'policy_net': { + 'ego_encoder': {'input_dim': 5, 'hidden_n': 0, 'hidden_dim':0, 'output_dim': 0}, + 'deepsets': { + 'input_dim': 6, + 'phi': { + 'hidden_n': ray_config['deepsets_phi_hidden_n'], + 'hidden_dim': ray_config['deepsets_phi_hidden_dim'] + }, + 'latent_dim': ray_config['deepsets_latent_dim'], + 'rho': { + 'hidden_n': ray_config['deepsets_rho_hidden_n'], + 'hidden_dim': ray_config['deepsets_rho_hidden_dim'] + }, + 'output_dim': ray_config['deepsets_output_dim'] + }, + 'path_encoder': {'input_dim': 40, 'hidden_n': 0, 'hidden_dim': 0, 'output_dim': 0}, + 'head': { + 'input_dim': 0, # computed in constructor + 'hidden_n': ray_config['head_hidden_n'], + 'hidden_dim': ray_config['head_hidden_dim'], + 'output_dim': 1, # number of outputs e.g. number of actions, or just one + 'final_activation': ray_config['head_final_activation'], + }, }, - 'path_encoder': {'input_dim': 40, 'hidden_n': 0, 'hidden_dim': 0, 'output_dim': 0}, - 'head': { - 'input_dim': 0, # computed in constructor - 'hidden_n': ray_config['head_hidden_n'], - 'hidden_dim': ray_config['head_hidden_dim'], - 'output_dim': 1, # number of outputs e.g. number of actions, or just one - 'final_activation': ray_config['head_final_activation'], + 'value_net': { + 'ego_encoder': {'input_dim': 5, 'hidden_n': 0, 'hidden_dim':0, 'output_dim': 0}, + 'deepsets': { + 'input_dim': 6, + 'phi': { + 'hidden_n': ray_config['deepsets_phi_hidden_n'], + 'hidden_dim': ray_config['deepsets_phi_hidden_dim'] + }, + 'latent_dim': ray_config['deepsets_latent_dim'], + 'rho': { + 'hidden_n': ray_config['deepsets_rho_hidden_n'], + 'hidden_dim': ray_config['deepsets_rho_hidden_dim'] + }, + 'output_dim': ray_config['deepsets_output_dim'] + }, + 'path_encoder': {'input_dim': 40, 'hidden_n': 0, 'hidden_dim': 0, 'output_dim': 0}, + 'action_dim': 1, + 'head': { + 'input_dim': 0, # computed in constructor + 'hidden_n': ray_config['head_hidden_n'], + 'hidden_dim': ray_config['head_hidden_dim'], + 'output_dim': 1, # number of outputs e.g. number of actions, or just one + 'final_activation': ray_config['head_final_activation'], + }, }, - 'optim': { + 'policy_optim': { 'optimizer':'adam', - 'lr':ray_config['lr'], - 'weight_decay':ray_config['weight_decay'] + 'lr':ray_config['policy_lr'], + 'weight_decay':ray_config['policy_weight_decay'] + }, + 'value_optim': { + 'optimizer':'adam', + 'lr':ray_config['value_lr'], + 'weight_decay':ray_config['value_weight_decay'] }, 'train_epochs': 40, 'train_batch_size': ray_config['train_batch_size'], - 'loss': ray_config['loss'], - + 'discount': ray_config['discount'], + 'clip_grad_norm': ray_config['clip_grad_norm'], } return config