Add value dice ray config
This commit is contained in:
@@ -90,6 +90,9 @@ def get_full_config(ray_config:dict, method:str)->dict:
|
|||||||
if method == 'bc':
|
if method == 'bc':
|
||||||
from src.bc import bc_config
|
from src.bc import bc_config
|
||||||
config = bc_config(ray_config)
|
config = bc_config(ray_config)
|
||||||
|
elif method == 'vd':
|
||||||
|
from src.value_dice import vd_config
|
||||||
|
config = vd_config(ray_config)
|
||||||
else:
|
else:
|
||||||
raise NotImplementedError
|
raise NotImplementedError
|
||||||
return config
|
return config
|
||||||
@@ -118,6 +121,25 @@ def get_ray_config(method:str)->dict:
|
|||||||
"head_hidden_dim": tune.lograndint(16,257),
|
"head_hidden_dim": tune.lograndint(16,257),
|
||||||
"head_final_activation": tune.choice(['sigmoid', None]),
|
"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:
|
else:
|
||||||
raise NotImplementedError
|
raise NotImplementedError
|
||||||
return ray_config
|
return ray_config
|
||||||
|
|||||||
@@ -11,39 +11,72 @@ from src.util.nn_training import optimizer_factory
|
|||||||
from tqdm import tqdm
|
from tqdm import tqdm
|
||||||
import json5
|
import json5
|
||||||
from ray import tune
|
from ray import tune
|
||||||
|
|
||||||
def vd_config(ray_config):
|
def vd_config(ray_config):
|
||||||
config = {
|
config = {
|
||||||
'ego_encoder': {'input_dim': 5, 'hidden_n': 0, 'hidden_dim':0, 'output_dim': 0},
|
'policy_net': {
|
||||||
'deepsets': {
|
'ego_encoder': {'input_dim': 5, 'hidden_n': 0, 'hidden_dim':0, 'output_dim': 0},
|
||||||
'input_dim': 6,
|
'deepsets': {
|
||||||
'phi': {
|
'input_dim': 6,
|
||||||
'hidden_n': ray_config['deepsets_phi_hidden_n'],
|
'phi': {
|
||||||
'hidden_dim': ray_config['deepsets_phi_hidden_dim']
|
'hidden_n': ray_config['deepsets_phi_hidden_n'],
|
||||||
},
|
'hidden_dim': ray_config['deepsets_phi_hidden_dim']
|
||||||
'latent_dim': ray_config['deepsets_latent_dim'],
|
},
|
||||||
'rho': {
|
'latent_dim': ray_config['deepsets_latent_dim'],
|
||||||
'hidden_n': ray_config['deepsets_rho_hidden_n'],
|
'rho': {
|
||||||
'hidden_dim': ray_config['deepsets_rho_hidden_dim']
|
'hidden_n': ray_config['deepsets_rho_hidden_n'],
|
||||||
},
|
'hidden_dim': ray_config['deepsets_rho_hidden_dim']
|
||||||
'output_dim': ray_config['deepsets_output_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},
|
'value_net': {
|
||||||
'head': {
|
'ego_encoder': {'input_dim': 5, 'hidden_n': 0, 'hidden_dim':0, 'output_dim': 0},
|
||||||
'input_dim': 0, # computed in constructor
|
'deepsets': {
|
||||||
'hidden_n': ray_config['head_hidden_n'],
|
'input_dim': 6,
|
||||||
'hidden_dim': ray_config['head_hidden_dim'],
|
'phi': {
|
||||||
'output_dim': 1, # number of outputs e.g. number of actions, or just one
|
'hidden_n': ray_config['deepsets_phi_hidden_n'],
|
||||||
'final_activation': ray_config['head_final_activation'],
|
'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',
|
'optimizer':'adam',
|
||||||
'lr':ray_config['lr'],
|
'lr':ray_config['policy_lr'],
|
||||||
'weight_decay':ray_config['weight_decay']
|
'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_epochs': 40,
|
||||||
'train_batch_size': ray_config['train_batch_size'],
|
'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
|
return config
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user