def build_model_tails(self, out_dim, out_layer_activation):
'''Build each model_tail. These are stored as Sequential models in model_tails'''
if not ps.is_list(out_layer_activation):
out_layer_activation = [out_layer_activation] * len(out_dim)
model_tails = nn.ModuleList()
if ps.is_empty(self.tail_hid_layers):
for out_d, out_activ in zip(out_dim, out_layer_activation):
tail = net_util.build_fc_model(
[self.body_hid_layers[-1], out_d], out_activ)
model_tails.append(tail)
else:
assert len(self.tail_hid_layers) == len(
out_dim
), 'Hydra tail hid_params inconsistent with number out dims'
for out_d, out_activ, hid_layers in zip(out_dim,
out_layer_activation,
self.tail_hid_layers):
dims = hid_layers
model_tail = net_util.build_fc_model(
dims, self.hid_layers_activation)
tail_out = net_util.build_fc_model([dims[-1], out_d],
out_activ)
model_tail.add_module(str(len(model_tail)), tail_out)
model_tails.append(model_tail)
return model_tails
def __init__(self, net_spec, in_dim, out_dim):
assert len(in_dim) == 3 # image shape (c,w,h)
nn.Module.__init__(self)
Net.__init__(self, net_spec, in_dim, out_dim)
# set default
util.set_attr(self, dict(
init_fn=None,
normalize=False,
batch_norm=False,
clip_grad_val=None,
loss_spec={'name': 'MSELoss'},
optim_spec={'name': 'Adam'},
lr_scheduler_spec=None,
update_type='replace',
update_frequency=1,
polyak_coef=0.0,
gpu=False,
))
util.set_attr(self, self.net_spec, [
'conv_hid_layers',
'fc_hid_layers',
'hid_layers_activation',
'init_fn',
'normalize',
'batch_norm',
'clip_grad_val',
'loss_spec',
'optim_spec',
'lr_scheduler_spec',
'update_type',
'update_frequency',
'polyak_coef',
'gpu',
])
# Guard against inappropriate algorithms and environments
assert isinstance(out_dim, int)
# conv body
self.conv_model = self.build_conv_layers(self.conv_hid_layers)
self.conv_out_dim = self.get_conv_output_size()
# fc body
if ps.is_empty(self.fc_hid_layers):
tail_in_dim = self.conv_out_dim
else:
# fc layer from flattened conv
self.fc_model = net_util.build_fc_model([self.conv_out_dim] + self.fc_hid_layers,
self.hid_layers_activation)
tail_in_dim = self.fc_hid_layers[-1]
# tails. avoid list for single-tail for compute speed
self.v = nn.Linear(tail_in_dim, 1) # state value
self.adv = nn.Linear(tail_in_dim, out_dim) # action dependent raw advantage
self.model_tails = nn.ModuleList(self.v, self.adv)
net_util.init_layers(self, self.init_fn)
self.loss_fn = net_util.get_loss_fn(self, self.loss_spec)
self.to(self.device)
self.train()
def build_model_heads(self, in_dim):
'''Build each model_head. These are stored as Sequential models in model_heads'''
assert len(self.head_hid_layers) == len(in_dim), 'Hydra head hid_params inconsistent with number in dims'
model_heads = nn.ModuleList()
for in_d, hid_layers in zip(in_dim, self.head_hid_layers):
dims = [in_d] + hid_layers
model_head = net_util.build_fc_model(dims, self.hid_layers_activation)
model_heads.append(model_head)
return model_heads
def __init__(self, net_spec, in_dim, out_dim):
nn.Module.__init__(self)
Net.__init__(self, net_spec, in_dim, out_dim)
# set default
util.set_attr(
self,
dict(
init_fn=None,
clip_grad_val=None,
loss_spec={'name': 'MSELoss'},
optim_spec={'name': 'Adam'},
lr_scheduler_spec=None,
update_type='replace',
update_frequency=1,
polyak_coef=0.0,
gpu=False,
))
util.set_attr(self, self.net_spec, [
'shared',
'hid_layers',
'hid_layers_activation',
'init_fn',
'clip_grad_val',
'loss_spec',
'optim_spec',
'lr_scheduler_spec',
'update_type',
'update_frequency',
'polyak_coef',
'gpu',
])
# Guard against inappropriate algorithms and environments
# Build model body
dims = [self.in_dim] + self.hid_layers
self.model_body = net_util.build_fc_model(dims,
self.hid_layers_activation)
# output layers
self.v = nn.Linear(dims[-1], 1) # state value
self.adv = nn.Linear(dims[-1],
out_dim) # action dependent raw advantage
net_util.init_layers(self, self.init_fn)
self.loss_fn = net_util.get_loss_fn(self, self.loss_spec)
self.to(self.device)
def __init__(self, net_spec, in_dim, out_dim):
'''
net_spec:
conv_hid_layers: list containing dimensions of the convolutional hidden layers, each is a list representing hid_layer = out_d, kernel, stride, padding, dilation.
Asssumed to all come before the flat layers.
Note: a convolutional layer should specify the in_channel, out_channels, kernel_size, stride (of kernel steps), padding, and dilation (spacing between kernel points) E.g. [3, 16, (5, 5), 1, 0, (2, 2)]
For more details, see http://pytorch.org/docs/master/nn.html#conv2d and https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md
fc_hid_layers: list of fc layers following the convolutional layers
hid_layers_activation: activation function for the hidden layers
out_layer_activation: activation function for the output layer, same shape as out_dim
init_fn: weight initialization function
normalize: whether to divide by 255.0 to normalize image input
batch_norm: whether to add batch normalization after each convolutional layer, excluding the input layer.
clip_grad_val: clip gradient norm if value is not None
loss_spec: measure of error between model predictions and correct outputs
optim_spec: parameters for initializing the optimizer
lr_scheduler_spec: Pytorch optim.lr_scheduler
update_type: method to update network weights: 'replace' or 'polyak'
update_frequency: how many total timesteps per update
polyak_coef: ratio of polyak weight update
gpu: whether to train using a GPU. Note this will only work if a GPU is available, othewise setting gpu=True does nothing
'''
assert len(in_dim) == 3 # image shape (c,w,h)
nn.Module.__init__(self)
super().__init__(net_spec, in_dim, out_dim)
# set default
util.set_attr(
self,
dict(
out_layer_activation=None,
init_fn=None,
normalize=False,
batch_norm=True,
clip_grad_val=None,
loss_spec={'name': 'MSELoss'},
optim_spec={'name': 'Adam'},
lr_scheduler_spec=None,
update_type='replace',
update_frequency=1,
polyak_coef=0.0,
gpu=False,
))
util.set_attr(self, self.net_spec, [
'conv_hid_layers',
'fc_hid_layers',
'hid_layers_activation',
'out_layer_activation',
'init_fn',
'normalize',
'batch_norm',
'clip_grad_val',
'loss_spec',
'optim_spec',
'lr_scheduler_spec',
'update_type',
'update_frequency',
'polyak_coef',
'gpu',
])
# conv body
self.conv_model = self.build_conv_layers(self.conv_hid_layers)
self.conv_out_dim = self.get_conv_output_size()
# fc body
if ps.is_empty(self.fc_hid_layers):
tail_in_dim = self.conv_out_dim
else:
# fc body from flattened conv
self.fc_model = net_util.build_fc_model([self.conv_out_dim] +
self.fc_hid_layers,
self.hid_layers_activation)
tail_in_dim = self.fc_hid_layers[-1]
# tails. avoid list for single-tail for compute speed
if ps.is_integer(self.out_dim):
self.model_tail = net_util.build_fc_model(
[tail_in_dim, self.out_dim], self.out_layer_activation)
else:
if not ps.is_list(self.out_layer_activation):
self.out_layer_activation = [self.out_layer_activation
] * len(out_dim)
assert len(self.out_layer_activation) == len(self.out_dim)
tails = []
for out_d, out_activ in zip(self.out_dim,
self.out_layer_activation):
tail = net_util.build_fc_model([tail_in_dim, out_d], out_activ)
tails.append(tail)
self.model_tails = nn.ModuleList(tails)
net_util.init_layers(self, self.init_fn)
self.loss_fn = net_util.get_loss_fn(self, self.loss_spec)
self.to(self.device)
self.train()
def __init__(self, net_spec, in_dim, out_dim):
'''
net_spec:
hid_layers: list containing dimensions of the hidden layers
hid_layers_activation: activation function for the hidden layers
out_layer_activation: activation function for the output layer, same shape as out_dim
init_fn: weight initialization function
clip_grad_val: clip gradient norm if value is not None
loss_spec: measure of error between model predictions and correct outputs
optim_spec: parameters for initializing the optimizer
lr_scheduler_spec: Pytorch optim.lr_scheduler
update_type: method to update network weights: 'replace' or 'polyak'
update_frequency: how many total timesteps per update
polyak_coef: ratio of polyak weight update
gpu: whether to train using a GPU. Note this will only work if a GPU is available, othewise setting gpu=True does nothing
'''
nn.Module.__init__(self)
super().__init__(net_spec, in_dim, out_dim)
# set default
util.set_attr(
self,
dict(
out_layer_activation=None,
init_fn=None,
clip_grad_val=None,
loss_spec={'name': 'MSELoss'},
optim_spec={'name': 'Adam'},
lr_scheduler_spec=None,
update_type='replace',
update_frequency=1,
polyak_coef=0.0,
gpu=False,
))
util.set_attr(self, self.net_spec, [
'shared',
'hid_layers',
'hid_layers_activation',
'out_layer_activation',
'init_fn',
'clip_grad_val',
'loss_spec',
'optim_spec',
'lr_scheduler_spec',
'update_type',
'update_frequency',
'polyak_coef',
'gpu',
])
dims = [self.in_dim] + self.hid_layers
self.model = net_util.build_fc_model(dims, self.hid_layers_activation)
# add last layer with no activation
# tails. avoid list for single-tail for compute speed
if ps.is_integer(self.out_dim):
self.model_tail = net_util.build_fc_model(
[dims[-1], self.out_dim], self.out_layer_activation)
else:
if not ps.is_list(self.out_layer_activation):
self.out_layer_activation = [self.out_layer_activation
] * len(out_dim)
assert len(self.out_layer_activation) == len(self.out_dim)
tails = []
for out_d, out_activ in zip(self.out_dim,
self.out_layer_activation):
tail = net_util.build_fc_model([dims[-1], out_d], out_activ)
tails.append(tail)
self.model_tails = nn.ModuleList(tails)
net_util.init_layers(self, self.init_fn)
self.loss_fn = net_util.get_loss_fn(self, self.loss_spec)
self.to(self.device)
self.train()
def __init__(self, net_spec, in_dim, out_dim):
'''
Multi state processing heads, single shared body, and multi action tails.
There is one state and action head per body/environment
Example:
env 1 state env 2 state
_______|______ _______|______
| head 1 | | head 2 |
|______________| |______________|
| |
|__________________|
________________|_______________
| Shared body |
|________________________________|
|
________|_______
| |
_______|______ ______|_______
| tail 1 | | tail 2 |
|______________| |______________|
| |
env 1 action env 2 action
'''
nn.Module.__init__(self)
super().__init__(net_spec, in_dim, out_dim)
# set default
util.set_attr(
self,
dict(
out_layer_activation=None,
init_fn=None,
clip_grad_val=None,
loss_spec={'name': 'MSELoss'},
optim_spec={'name': 'Adam'},
lr_scheduler_spec=None,
update_type='replace',
update_frequency=1,
polyak_coef=0.0,
gpu=False,
))
util.set_attr(self, self.net_spec, [
'hid_layers',
'hid_layers_activation',
'out_layer_activation',
'init_fn',
'clip_grad_val',
'loss_spec',
'optim_spec',
'lr_scheduler_spec',
'update_type',
'update_frequency',
'polyak_coef',
'gpu',
])
assert len(
self.hid_layers
) == 3, 'Your hidden layers must specify [*heads], [body], [*tails]. If not, use MLPNet'
assert isinstance(self.in_dim,
list), 'Hydra network needs in_dim as list'
assert isinstance(self.out_dim,
list), 'Hydra network needs out_dim as list'
self.head_hid_layers = self.hid_layers[0]
self.body_hid_layers = self.hid_layers[1]
self.tail_hid_layers = self.hid_layers[2]
if len(self.head_hid_layers) == 1:
self.head_hid_layers = self.head_hid_layers * len(self.in_dim)
if len(self.tail_hid_layers) == 1:
self.tail_hid_layers = self.tail_hid_layers * len(self.out_dim)
self.model_heads = self.build_model_heads(in_dim)
heads_out_dim = np.sum(
[head_hid_layers[-1] for head_hid_layers in self.head_hid_layers])
dims = [heads_out_dim] + self.body_hid_layers
self.model_body = net_util.build_fc_model(dims,
self.hid_layers_activation)
self.model_tails = self.build_model_tails(self.out_dim,
self.out_layer_activation)
net_util.init_layers(self, self.init_fn)
self.loss_fn = net_util.get_loss_fn(self, self.loss_spec)
self.to(self.device)
self.train()
def __init__(self, net_spec, in_dim, out_dim):
'''
net_spec:
cell_type: any of RNN, LSTM, GRU
fc_hid_layers: list of fc layers preceeding the RNN layers
hid_layers_activation: activation function for the fc hidden layers
out_layer_activation: activation function for the output layer, same shape as out_dim
rnn_hidden_size: rnn hidden_size
rnn_num_layers: number of recurrent layers
bidirectional: if RNN should be bidirectional
seq_len: length of the history of being passed to the net
init_fn: weight initialization function
clip_grad_val: clip gradient norm if value is not None
loss_spec: measure of error between model predictions and correct outputs
optim_spec: parameters for initializing the optimizer
lr_scheduler_spec: Pytorch optim.lr_scheduler
update_type: method to update network weights: 'replace' or 'polyak'
update_frequency: how many total timesteps per update
polyak_coef: ratio of polyak weight update
gpu: whether to train using a GPU. Note this will only work if a GPU is available, othewise setting gpu=True does nothing
'''
nn.Module.__init__(self)
super().__init__(net_spec, in_dim, out_dim)
# set default
util.set_attr(
self,
dict(
out_layer_activation=None,
cell_type='GRU',
rnn_num_layers=1,
bidirectional=False,
init_fn=None,
clip_grad_val=None,
loss_spec={'name': 'MSELoss'},
optim_spec={'name': 'Adam'},
lr_scheduler_spec=None,
update_type='replace',
update_frequency=1,
polyak_coef=0.0,
gpu=False,
))
util.set_attr(self, self.net_spec, [
'cell_type',
'fc_hid_layers',
'hid_layers_activation',
'out_layer_activation',
'rnn_hidden_size',
'rnn_num_layers',
'bidirectional',
'seq_len',
'init_fn',
'clip_grad_val',
'loss_spec',
'optim_spec',
'lr_scheduler_spec',
'update_type',
'update_frequency',
'polyak_coef',
'gpu',
])
# restore proper in_dim from env stacked state_dim (stack_len, *raw_state_dim)
self.in_dim = in_dim[1:] if len(in_dim) > 2 else in_dim[1]
# fc body: state processing model
if ps.is_empty(self.fc_hid_layers):
self.rnn_input_dim = self.in_dim
else:
fc_dims = [self.in_dim] + self.fc_hid_layers
self.fc_model = net_util.build_fc_model(fc_dims,
self.hid_layers_activation)
self.rnn_input_dim = fc_dims[-1]
# RNN model
self.rnn_model = getattr(nn, net_util.get_nn_name(self.cell_type))(
input_size=self.rnn_input_dim,
hidden_size=self.rnn_hidden_size,
num_layers=self.rnn_num_layers,
batch_first=True,
bidirectional=self.bidirectional)
# tails. avoid list for single-tail for compute speed
if ps.is_integer(self.out_dim):
self.model_tail = net_util.build_fc_model(
[self.rnn_hidden_size, self.out_dim],
self.out_layer_activation)
else:
if not ps.is_list(self.out_layer_activation):
self.out_layer_activation = [self.out_layer_activation
] * len(out_dim)
assert len(self.out_layer_activation) == len(self.out_dim)
tails = []
for out_d, out_activ in zip(self.out_dim,
self.out_layer_activation):
tail = net_util.build_fc_model([self.rnn_hidden_size, out_d],
out_activ)
tails.append(tail)
self.model_tails = nn.ModuleList(tails)
net_util.init_layers(self, self.init_fn)
self.loss_fn = net_util.get_loss_fn(self, self.loss_spec)
self.to(self.device)
self.train()