def __init__(self,
input_shape: Tuple,
output_size: int,
rnn_type: str = 'gru',
rnn_size: int = 256,
hidden_sizes: [List, Tuple] = None,
baselines_init: bool = True,
layer_norm: bool = False
):
super().__init__()
self._obs_dim = 2
self.rnn_is_lstm = rnn_type != 'gru'
input_size = int(np.prod(input_shape))
rnn_class = get_rnn_class(rnn_type, layer_norm)
self.body_pi = MlpModel(input_size, hidden_sizes, None, nn.ReLU, None)
self.body_v = MlpModel(input_size, hidden_sizes, None, nn.ReLU, None)
self.rnn_pi = rnn_class(self.body_pi.output_size + output_size + 1, rnn_size) # Concat action, reward
self.rnn_v = rnn_class(self.body_v.output_size + output_size + 1, rnn_size)
self.pi = nn.Sequential(nn.ReLU(), nn.Linear(rnn_size, output_size), nn.Softmax(-1)) # Need to activate after lstm
self.v = nn.Sequential(nn.ReLU(), nn.Linear(rnn_size, 1))
if baselines_init:
self.body_pi.apply(apply_init); self.body_v.apply(apply_init)
self.rnn_pi.apply(apply_init); self.rnn_v.apply(apply_init)
self.pi.apply(partial(apply_init, O_INIT_VALUES['pi']))
self.v.apply(partial(apply_init, O_INIT_VALUES['v']))
self.body_pi, self.body_v, self.pi, self.v = tscr(self.body_pi), tscr(self.body_v), tscr(self.pi), tscr(self.v)
def __init__(
self,
observation_shape,
action_size,
hidden_sizes=None, # None for default (see below).
hidden_nonlinearity=torch.nn.Tanh, # Module form.
mu_nonlinearity=torch.nn.Tanh, # Module form.
init_log_std=0.,
):
"""Instantiate neural net modules according to inputs."""
super().__init__()
self._obs_ndim = len(observation_shape)
input_size = int(np.prod(observation_shape))
hidden_sizes = hidden_sizes or [64, 64]
mu_mlp = MlpModel(
input_size=input_size,
hidden_sizes=hidden_sizes,
output_size=action_size,
nonlinearity=hidden_nonlinearity,
)
if mu_nonlinearity is not None:
self.mu = torch.nn.Sequential(mu_mlp, mu_nonlinearity())
else:
self.mu = mu_mlp
self.v = MlpModel(
input_size=input_size,
hidden_sizes=hidden_sizes,
output_size=1,
nonlinearity=hidden_nonlinearity,
)
self.log_std = torch.nn.Parameter(init_log_std *
torch.ones(action_size))
def __init__(
self,
observation_shape,
hidden_sizes,
action_size,
n_tile=20,
):
super().__init__()
self._obs_ndim = 1
self._n_tile = n_tile
input_dim = int(np.sum(observation_shape))
self._action_size = action_size
self.mlp_loc = MlpModel(
input_size=input_dim,
hidden_sizes=hidden_sizes,
output_size=4
)
self.mlp_delta = MlpModel(
input_size=input_dim + 4 * n_tile,
hidden_sizes=hidden_sizes,
output_size=3 * 2,
)
self.delta_distribution = Gaussian(
dim=3,
squash=True,
min_std=np.exp(MIN_LOG_STD),
max_std=np.exp(MAX_LOG_STD),
)
self.cat_distribution = Categorical(4)
self._counter = 0
def __init__(
self,
observation_shape,
hidden_sizes,
action_size,
):
super().__init__()
assert hasattr(observation_shape, 'state'), "mcp model requires observation dict to contain state attribute"
assert hasattr(observation_shape, 'goal'), "mcp model requires observation to contain goal attribute"
self.height, self.width = observation_shape.goal
# self.conv = Conv2dModel(
# in_channels=1,
# channels=[8, 20],
# kernel_sizes=[5, 4],
# strides=[3, 3],
# paddings=[1, 1],
# )
# conv_out_size = self.conv.conv_out_size(self.height, self.width)
self.robot_state_mlp = MlpModel(
input_size=observation_shape.state[0],
hidden_sizes=[512],
output_size=256
)
self.mlp = MlpModel(
input_size=256 + observation_shape.goal.relative_target[0] + action_size,
hidden_sizes=[256],
output_size=1
)
def __init__(
self,
observation_shape,
action_size,
fc_sizes=512,
dueling=False,
use_maxpool=False,
channels=None, # None uses default.
kernel_sizes=None,
strides=None,
paddings=None,
linear_value_output=True):
"""Instantiates the neural network according to arguments; network defaults
stored within this method."""
super().__init__()
self.dueling = dueling
c, h, w = observation_shape
self.conv = Conv2dModel(
in_channels=c,
channels=channels or [32, 64, 64],
kernel_sizes=kernel_sizes or [3, 2, 2],
strides=strides or [2, 1, 1],
paddings=paddings or [0, 0, 0],
use_maxpool=use_maxpool,
)
conv_out_size = self.conv.conv_out_size(h, w)
self.pi_head = MlpModel(conv_out_size, [
256,
], action_size)
self.value_head = MlpModel(conv_out_size, [
256,
], 1 if linear_value_output else None)
def __init__(
self,
observation_shape,
hidden_sizes,
action_size,
n_tile=50,
loc_size=2,
delta_size=3,
):
super().__init__()
self._obs_ndim = 1
input_dim = int(np.sum(observation_shape))
self._n_tile = n_tile
self._loc_size = loc_size
self._delta_size = delta_size
# self._obs_ndim = len(observation_shape)
# input_dim = int(np.prod(observation_shape))
assert action_size == loc_size + delta_size # First 2 (location), then 3 (action)
self._action_size = action_size
self.mlp_loc = MlpModel(
input_size=input_dim,
hidden_sizes=hidden_sizes,
output_size=loc_size * 2
)
self.mlp_delta = MlpModel(
input_size=input_dim + loc_size * n_tile,
hidden_sizes=hidden_sizes,
output_size=delta_size * 2,
)
self._counter = 0
def __init__(
self,
observation_shape,
hidden_sizes,
action_size,
):
super().__init__()
assert hasattr(observation_shape, 'state'), "vision model requires observation dict to contain state attribute"
assert hasattr(observation_shape, 'goal'), "vision model requires observation to contain goal attribute"
assert hasattr(observation_shape.goal, 'camera'), 'vision model requires camera observation'
self.height, self.width = observation_shape.goal.camera
robot_state_shape = observation_shape.state[0]
self.conv = Conv2dModel(
in_channels=1,
channels=[9, 18],
kernel_sizes=[3, 3],
strides=[2, 2],
paddings=[1, 1],
)
conv_out_size = self.conv.conv_out_size(self.height, self.width)
robot_state_out = 256
self.robot_state_mlp = MlpModel(
input_size=robot_state_shape,
hidden_sizes=[],
output_size=robot_state_out
)
self.mu_head = MlpModel(
input_size=robot_state_out + conv_out_size,
hidden_sizes=[256, ],
output_size=action_size
)
init_log_std = 0.
self.log_std = torch.nn.Parameter(init_log_std * torch.ones(action_size))
def __init__(self,
input_size,
hidden_sizes,
output_size,
n_atoms,
grad_scale=2**(-1 / 2),
fc_1_V=None):
"""
Dueling distributional C51 head copied from rlpyt
"""
super().__init__()
if isinstance(hidden_sizes, int):
hidden_sizes = [hidden_sizes]
self.advantage_hidden = MlpModel(input_size, hidden_sizes)
self.advantage_out = torch.nn.Linear(hidden_sizes[-1],
output_size * n_atoms,
bias=False)
self.advantage_bias = torch.nn.Parameter(torch.zeros(n_atoms))
if fc_1_V is None:
self.value = MlpModel(input_size,
hidden_sizes,
output_size=n_atoms)
else:
self.value = nn.Sequential(
*[fc_1_V,
nn.ReLU(),
nn.Linear(hidden_sizes[0], n_atoms)])
self._grad_scale = grad_scale
self._output_size = output_size
self._n_atoms = n_atoms
def __init__(self,
input_classes: int,
output_size: int,
rnn_type: str = 'gru',
rnn_size: int = 256,
hidden_sizes: [List, Tuple] = None,
baselines_init: bool = True,
layer_norm: bool = False,
prev_action: int = 3,
prev_reward: int = 3,
):
super().__init__()
self._obs_dim = 0
self.rnn_is_lstm = rnn_type != 'gru'
self.preprocessor = tscr(OneHotLayer(input_classes))
rnn_class = get_rnn_class(rnn_type, layer_norm)
self.body_pi = MlpModel(input_classes, hidden_sizes, None, nn.ReLU, None)
self.body_v = MlpModel(input_classes, hidden_sizes, None, nn.ReLU, None)
rnn_input_size_pi = self.body_pi.output_size + (prev_action in [1,3]) * output_size + (prev_reward in [1,3])
rnn_input_size_v = self.body_v.output_size + (prev_action in [2,3]) * output_size + (prev_reward in [2,3])
self.rnn_pi = rnn_class(rnn_input_size_pi, rnn_size) # Concat action, reward
self.rnn_v = rnn_class(rnn_input_size_v, rnn_size)
self.pi = nn.Sequential(nn.ReLU(), nn.Linear(rnn_size, output_size), nn.Softmax(-1)) # Need to activate after lstm
self.v = nn.Sequential(nn.ReLU(), nn.Linear(rnn_size, 1))
if baselines_init:
self.body_pi.apply(apply_init); self.body_v.apply(apply_init)
self.rnn_pi.apply(apply_init); self.rnn_v.apply(apply_init)
self.pi.apply(partial(apply_init, O_INIT_VALUES['pi']))
self.v.apply(partial(apply_init, O_INIT_VALUES['v']))
self.body_pi, self.body_v, self.pi, self.v = tscr(self.body_pi), tscr(self.body_v), tscr(self.pi), tscr(self.v)
self.p_a = prev_action
self.p_r = prev_reward
def __init__(self,
input_classes: int,
output_size: int,
rnn_type: str = 'gru',
rnn_size: int = 256,
hidden_sizes: [List, Tuple] = None,
baselines_init: bool = True,
layer_norm: bool = False,
prev_action: int = 2,
prev_reward: int = 2,
):
super().__init__()
self._obs_dim = 0
self.rnn_is_lstm = rnn_type != 'gru'
self.preprocessor = tscr(OneHotLayer(input_classes))
rnn_class = get_rnn_class(rnn_type, layer_norm)
rnn_input_size = input_classes
if prev_action: rnn_input_size += output_size # Use previous action as input
if prev_reward: rnn_input_size += 1 # Use previous reward as input
self.rnn = rnn_class(rnn_input_size, rnn_size) # Concat action, reward
pi_inits = (O_INIT_VALUES['base'], O_INIT_VALUES['pi']) if baselines_init else None
v_inits = (O_INIT_VALUES['base'], O_INIT_VALUES['v']) if baselines_init else None
self.pi = nn.Sequential(MlpModel(rnn_size, hidden_sizes, output_size, nn.ReLU, pi_inits), nn.Softmax(-1))
self.v = nn.Sequential(MlpModel(rnn_size, hidden_sizes, 1, nn.ReLU, v_inits))
if baselines_init:
self.rnn.apply(apply_init)
self.pi, self.v = tscr(self.pi), tscr(self.v)
self.p_a = prev_action > 0
self.p_r = prev_reward > 0
def __init__(
self,
observation_shape,
output_size,
fc_size=512, # Between mlp and lstm.
lstm_size=512,
head_size=256,
dueling=False,
):
"""Instantiates the neural network according to arguments; network defaults
stored within this method."""
super().__init__()
self._obs_n_dim = len(observation_shape)
self.dueling = dueling
self.mlp = MlpModel(
observation_shape,
[256],
output_size=fc_size,
nonlinearity=torch.nn.Tanh # Match spinningup
)
self.lstm = torch.nn.LSTM(fc_size + output_size + 1, lstm_size)
if dueling:
self.head = DuelingHeadModel(lstm_size, head_size, output_size)
else:
self.head = MlpModel(lstm_size, head_size, output_size=output_size)
def __init__(
self,
observation_shape,
action_size,
linear_value_output=True,
sequence_length=64,
seperate_value_network=True,
size='medium',
):
super().__init__()
self.state_size = np.prod(observation_shape.state)
self.action_size = action_size
self.sequence_length = sequence_length
self.transformer_dim = SIZES[size]['dim']
self.depth = SIZES[size]['depth']
self.cmem_ratio = SIZES[size]['cmem_ratio']
self.cmem_length = self.sequence_length // self.cmem_ratio
memory_layers = range(1, self.depth + 1)
self.transformer = CompressiveTransformerPyTorch(
num_tokens=20000,
emb_dim=self.
state_size, # embedding dimensions, embedding factorization from Albert paper
dim=self.transformer_dim,
heads=SIZES[size]['num_heads'],
depth=self.depth,
seq_len=self.sequence_length,
mem_len=self.sequence_length, # memory length
# cmem_len=self.cmem_length, # compressed memory buffer length
# cmem_ratio=self.cmem_ratio, # compressed memory ratio, 4 was recommended in paper
reconstruction_loss_weight=
1, # weight to place on compressed memory reconstruction loss
gru_gated_residual=True,
# whether to gate the residual intersection, from 'Stabilizing Transformer for RL' paper
memory_layers=memory_layers,
)
self.transformer.token_emb = torch.nn.Identity(
) # don't use token embedding in compressive transforrmer
self.transformer.to_logits = torch.nn.Identity()
# self.input_layer_norm = torch.nn.LayerNorm(self.state_size)
self.input_layer_norm = torch.nn.Identity()
# self.output_layer_norm = torch.nn.LayerNorm(self.transformer_dim)
self.output_layer_norm = torch.nn.Identity()
self.softplus = torch.nn.Softplus()
self.pi_head = MlpModel(input_size=self.transformer_dim,
hidden_sizes=[
256,
],
output_size=2 * action_size)
self.value_head = MlpModel(
input_size=self.transformer_dim,
hidden_sizes=[
256,
],
output_size=1 if linear_value_output else None)
self.mask = torch.ones((self.sequence_length, self.sequence_length),
dtype=torch.int8).triu()
def __init__(
self,
observation_shape,
action_size,
fc_sizes=32,
use_maxpool=False,
channels=None, # None uses default.
kernel_sizes=None,
strides=None,
paddings=None,
hidden_nonlinearity=torch.nn.Tanh, # Module form.
# mu_nonlinearity=torch.nn.Tanh, # Module form.
mu_nonlinearity=None,
init_log_std=0.,
):
"""Instantiate neural net module according to inputs."""
super().__init__()
self.conv = Conv2dHeadModel(
image_shape=observation_shape,
channels=channels or [16, 32],
kernel_sizes=kernel_sizes or [8, 4],
strides=strides or [4, 2],
paddings=paddings or [0, 1],
use_maxpool=use_maxpool,
hidden_sizes=fc_sizes, # Applies nonlinearity at end.
)
mu_mlp = MlpModel(
input_size=self.conv.output_size,
hidden_sizes=fc_sizes,
output_size=action_size,
nonlinearity=hidden_nonlinearity,
)
# print(self.conv.output_size)
# print('Num of encoder parameters: %d' % sum(p.numel() for p in self.conv.parameters() if p.requires_grad))
# print('Num of encoder parameters: %d' % sum(p.numel() for p in mu_mlp.parameters() if p.requires_grad))
if mu_nonlinearity is not None:
self.mu = torch.nn.Sequential(mu_mlp, mu_nonlinearity())
else:
self.mu = mu_mlp
self.v = MlpModel(
input_size=self.conv.output_size,
hidden_sizes=fc_sizes,
output_size=1,
nonlinearity=hidden_nonlinearity,
)
self.lstm = torch.nn.LSTM(mlp_output_size + action_size + 1, lstm_size)
self.head = torch.nn.Linear(lstm_size, action_size * 2 + 1)
self.log_std = torch.nn.Parameter(init_log_std *
torch.ones(action_size))
def __init__(self, observation_shape, action_size,
policy_hidden_sizes=None, policy_hidden_nonlinearity=torch.nn.Tanh,
value_hidden_sizes=None, value_hidden_nonlinearity=torch.nn.Tanh, ):
super().__init__()
self._obs_ndim = len(observation_shape)
input_size = int(np.prod(observation_shape))
policy_hidden_sizes = [400, 300] if policy_hidden_sizes is None else policy_hidden_sizes
value_hidden_sizes = [400, 300] if value_hidden_sizes is None else value_hidden_sizes
self.pi = MlpModel(input_size=input_size, hidden_sizes=policy_hidden_sizes, output_size=action_size,
nonlinearity=policy_hidden_nonlinearity)
self.v = MlpModel(input_size=input_size, hidden_sizes=value_hidden_sizes, output_size=1,
nonlinearity=value_hidden_nonlinearity, )
def __init__(
self,
observation_shape,
action_size,
hidden_sizes=None, # None for default (see below).
hidden_nonlinearity=torch.nn.Tanh, # Module form.
mu_nonlinearity=torch.nn.Tanh, # Module form.
init_log_std=0.,
pooling="average",
):
super().__init__()
self._obs_ndim = len(observation_shape)
self._n_pop = observation_shape[0]
input_size = int(observation_shape[-1])
output_size = int(action_size[-1])
hidden_sizes = hidden_sizes or [64]
# import pdb; pdb.set_trace()
self.encoder = MlpModel(input_size=input_size,
hidden_sizes=hidden_sizes * 2,
output_size=None,
nonlinearity=hidden_nonlinearity)
self.pooling = pooling
input_size = 64
if self.pooling is not None:
input_size *= 2
mu_mlp = MlpModel(
input_size=input_size,
hidden_sizes=hidden_sizes,
output_size=output_size,
nonlinearity=hidden_nonlinearity,
)
if mu_nonlinearity is not None:
self.mu = torch.nn.Sequential(mu_mlp, mu_nonlinearity())
else:
self.mu = mu_mlp
self.v = MlpModel(
input_size=input_size,
hidden_sizes=hidden_sizes,
output_size=1,
nonlinearity=hidden_nonlinearity,
)
self.log_std = torch.nn.Parameter(init_log_std *
torch.ones(action_size))
def __init__(
self,
image_shape,
output_size,
fc_size=512, # Between conv and lstm.
lstm_size=512,
head_size=512,
dueling=False,
use_maxpool=False,
channels=None, # None uses default.
kernel_sizes=None,
strides=None,
paddings=None,
):
super().__init__()
self.dueling = dueling
self.conv = Conv2dHeadModel(
image_shape=image_shape,
channels=channels or [32, 64, 64],
kernel_sizes=kernel_sizes or [8, 4, 3],
strides=strides or [4, 2, 1],
paddings=paddings or [0, 1, 1],
use_maxpool=use_maxpool,
hidden_sizes=fc_size,
)
self.lstm = torch.nn.LSTM(self.conv.output_size + output_size + 1,
lstm_size)
if dueling:
self.head = DuelingHeadModel(lstm_size, head_size, output_size)
else:
self.head = MlpModel(lstm_size, head_size, output_size=output_size)
def __init__(
self,
image_shape,
output_size,
fc_sizes=512,
dueling=False,
use_maxpool=False,
channels=None, # None uses default.
kernel_sizes=None,
strides=None,
paddings=None,
):
"""Instantiates the neural network according to arguments; network defaults
stored within this method."""
super().__init__()
self.dueling = dueling
c, h, w = image_shape
self.conv = Conv2dModel(
in_channels=c,
channels=channels or [32, 64, 64],
kernel_sizes=kernel_sizes or [8, 4, 3],
strides=strides or [4, 2, 1],
paddings=paddings or [0, 1, 1],
use_maxpool=use_maxpool,
)
conv_out_size = self.conv.conv_out_size(h, w)
if dueling:
self.head = DuelingHeadModel(conv_out_size, fc_sizes, output_size)
else:
self.head = MlpModel(conv_out_size, fc_sizes, output_size)
def __init__(
self,
image_shape,
channels,
kernel_sizes,
strides,
hidden_sizes,
output_size=None, # if None: nonlinearity applied to output.
paddings=None,
nonlinearity=torch.nn.ReLU,
use_maxpool=False,
):
super().__init__()
c, h, w = image_shape
self.conv = Conv2dModel(
in_channels=c,
channels=channels,
kernel_sizes=kernel_sizes,
strides=strides,
paddings=paddings,
nonlinearity=nonlinearity,
use_maxpool=use_maxpool,
)
conv_out_size = self.conv.conv_out_size(h, w)
if hidden_sizes or output_size:
self.head = MlpModel(conv_out_size, hidden_sizes,
output_size=output_size, nonlinearity=nonlinearity)
if output_size is not None:
self._output_size = output_size
else:
self._output_size = (hidden_sizes if
isinstance(hidden_sizes, int) else hidden_sizes[-1])
else:
self.head = lambda x: x
self._output_size = conv_out_size
def __init__(
self,
image_shape,
output_size,
fc_sizes=512,
dueling=False,
use_maxpool=False,
channels=None, # None uses default.
kernel_sizes=None,
strides=None,
paddings=None,
):
super().__init__()
self.dueling = dueling
c, h, w = image_shape
self.conv = Conv2dModel(
in_channels=c,
channels=channels or [32, 64, 64],
kernel_sizes=kernel_sizes or [8, 4, 3],
strides=strides or [4, 2, 1],
paddings=paddings or [0, 1, 1],
use_maxpool=use_maxpool,
)
conv_out_size = self.conv.conv_out_size(h, w)
if dueling:
self.head = DuelingHeadModel(conv_out_size, fc_sizes, output_size)
else:
self.head = MlpModel(conv_out_size, fc_sizes, output_size)
def __init__(
self,
observation_shape,
action_size,
hidden_sizes=None, # None for default (see below).
lstm_size=256,
nonlinearity=torch.nn.ReLU,
normalize_observation=False,
norm_obs_clip=10,
norm_obs_var_clip=1e-6,
):
super().__init__()
self._obs_n_dim = len(observation_shape)
self._action_size = action_size
hidden_sizes = hidden_sizes or [256, 256]
mlp_input_size = int(np.prod(observation_shape))
self.mlp = MlpModel(
input_size=mlp_input_size,
hidden_sizes=hidden_sizes,
output_size=None,
nonlinearity=nonlinearity,
)
mlp_output_size = hidden_sizes[-1] if hidden_sizes else mlp_input_size
self.lstm = torch.nn.LSTM(mlp_output_size + action_size + 1, lstm_size)
self.head = torch.nn.Linear(lstm_size, action_size * 2 + 1)
if normalize_observation:
self.obs_rms = RunningMeanStdModel(observation_shape)
self.norm_obs_clip = norm_obs_clip
self.norm_obs_var_clip = norm_obs_var_clip
self.normalize_observation = normalize_observation
def __init__(
self,
observation_shape,
action_size,
hidden_sizes=[64,64], # mlp after lstm
fc_sizes=128, # Between conv and lstm
lstm_size=64,
channels=[8, 16],
kernel_sizes=[8, 4],
strides=[4, 2],
paddings=[0, 1],
use_maxpool=False,
):
"""Instantiate neural net according to inputs."""
super().__init__()
self._obs_ndim = len(observation_shape)
self.conv = Conv2dHeadModel(
image_shape=observation_shape,
channels=channels,
kernel_sizes=kernel_sizes,
strides=strides,
paddings=paddings,
use_maxpool=use_maxpool,
hidden_sizes=fc_sizes, # Applies nonlinearity at end.
) # image -> conv (ReLU) -> linear (fc_sizes) - > ReLU
self.lstm = torch.nn.LSTM(self.conv.output_size + action_size + 1 + action_size, lstm_size) # Input to LSTM: conv_output + prev_action + prev_reward + action
self.mlp = MlpModel(
input_size=lstm_size,
hidden_sizes=hidden_sizes,
output_size=1,
)
def __init__(
self,
image_shape,
latent_size,
channels,
kernel_sizes,
strides,
paddings=None,
hidden_sizes=None, # usually None; NOT the same as anchor MLP
kiaming_init=True,
):
super().__init__()
c, h, w = image_shape
self.conv = Conv2dModel(
in_channels=c,
channels=channels,
kernel_sizes=kernel_sizes,
strides=strides,
paddings=paddings,
use_maxpool=False,
)
self._output_size = self.conv.conv_out_size(h, w)
self._output_shape = self.conv.conv_out_shape(h, w)
self.head = MlpModel(
input_size=self._output_size,
hidden_sizes=hidden_sizes,
output_size=latent_size,
)
if kiaming_init:
self.apply(weight_init)
def __init__(
self,
image_shape,
latent_size,
use_fourth_layer=True,
skip_connections=True,
hidden_sizes=None,
kiaming_init=True,
):
super().__init__()
c, h, w = image_shape
self.conv = DmlabConv2dModel(
in_channels=c,
use_fourth_layer=True,
skip_connections=skip_connections,
use_maxpool=False,
)
self._output_size = self.conv.output_size(h, w)
self._output_shape = self.conv.output_shape(h, w)
self.head = MlpModel( # gets to z_t, not necessarily c_t
input_size=self._output_size,
hidden_sizes=hidden_sizes,
output_size=latent_size,
)
if kiaming_init:
self.apply(weight_init)
def __init__(
self,
image_shape,
output_size,
fc_size=512, # Between conv and lstm.
lstm_size=512,
head_size=512,
dueling=False,
use_maxpool=False,
channels=None, # None uses default.
kernel_sizes=None,
strides=None,
paddings=None,
):
"""Instantiates the neural network according to arguments; network defaults
stored within this method."""
super().__init__()
self.dueling = dueling
self.conv = Conv2dHeadModel(
image_shape=image_shape,
channels=channels or [32, 64, 64],
kernel_sizes=kernel_sizes or [8, 4, 3],
strides=strides or [4, 2, 1],
paddings=paddings or [0, 1, 1],
use_maxpool=use_maxpool,
hidden_sizes=fc_size, # ReLU applied here (Steven).
)
self.lstm = torch.nn.LSTM(self.conv.output_size + output_size + 1,
lstm_size)
if dueling:
self.head = DuelingHeadModel(lstm_size, head_size, output_size)
else:
self.head = MlpModel(lstm_size, head_size, output_size=output_size)
def __init__(
self,
observation_shape,
output_size,
hidden_sizes=None, # None for default (see below).
lstm_size=256,
nonlinearity=torch.nn.ReLU,
normalize_observation=False,
norm_obs_clip=10,
norm_obs_var_clip=1e-6,
):
"""Instantiate neural net module according to inputs."""
super().__init__()
self._obs_n_dim = len(observation_shape)
hidden_sizes = hidden_sizes or [256, 256]
mlp_input_size = int(np.prod(observation_shape))
self.mlp = MlpModel(
input_size=mlp_input_size,
hidden_sizes=hidden_sizes,
output_size=None,
nonlinearity=nonlinearity,
)
mlp_output_size = hidden_sizes[-1] if hidden_sizes else mlp_input_size
self.lstm = torch.nn.LSTM(mlp_output_size + output_size + 1, lstm_size)
self.pi = torch.nn.Linear(lstm_size, output_size)
self.value = torch.nn.Linear(lstm_size, 1)
if normalize_observation:
self.obs_rms = RunningMeanStdModel(observation_shape)
self.norm_obs_clip = norm_obs_clip
self.norm_obs_var_clip = norm_obs_var_clip
self.normalize_observation = normalize_observation
def __init__(
self,
input_size,
hidden_sizes,
output_size,
grad_scale=2 ** (-1 / 2),
):
super().__init__()
if isinstance(hidden_sizes, int):
hidden_sizes = [hidden_sizes]
self.advantage_hidden = MlpModel(input_size, hidden_sizes)
self.advantage_out = torch.nn.Linear(hidden_sizes[-1], output_size,
bias=False)
self.advantage_bias = torch.nn.Parameter(torch.zeros(1))
self.value = MlpModel(input_size, hidden_sizes, output_size=1)
self._grad_scale = grad_scale
def __init__(
self,
image_shape,
latent_size,
channels=None,
kernel_sizes=None,
strides=None,
paddings=None,
hidden_sizes=None,
kiaming_init=True,
):
super().__init__()
c, h, w = image_shape
self.conv = Conv2dStdimModel(
in_channels=c,
channels=channels or [32, 64, 64],
kernel_sizes=kernel_sizes or [8, 4, 3],
strides=strides or [4, 2, 1],
paddings=paddings,
use_maxpool=False,
)
self._output_size = self.conv.conv_out_size(h, w)
self._output_shape = self.conv.conv_out_shape(h, w)
self._conv_layer_shapes = self.conv.conv_layer_shapes(h, w)
self.head = MlpModel(
input_size=self._output_size,
hidden_sizes=hidden_sizes,
output_size=latent_size,
)
if kiaming_init:
self.apply(weight_init)
def __init__(
self,
observation_shape,
action_size,
hidden_sizes=[64, 64], # mlp after lstm
fc_sizes=64, # Between conv and lstm
channels=None,
kernel_sizes=None,
strides=None,
paddings=None,
use_maxpool=False,
):
"""Instantiate neural net according to inputs."""
super().__init__()
self._obs_ndim = len(observation_shape)
self.conv = Conv2dHeadModel(
image_shape=observation_shape,
channels=channels or [4, 8],
kernel_sizes=kernel_sizes or [8, 4],
strides=strides or [4, 2],
paddings=paddings or [0, 1],
use_maxpool=use_maxpool,
hidden_sizes=fc_sizes, # Applies nonlinearity at end.
) # image -> conv (ReLU) -> linear (fc_sizes) - > ReLU
self.mlp = MlpModel(
input_size=self.conv.output_size + action_size,
hidden_sizes=hidden_sizes,
output_size=1,
)
def __init__(
self,
observation_shape,
hidden_sizes,
action_size,
all_corners=False
):
super().__init__()
self._obs_ndim = 1
self._all_corners = all_corners
input_dim = int(np.sum(observation_shape))
print('all corners', self._all_corners)
delta_dim = 12 if all_corners else 3
self._delta_dim = delta_dim
self.mlp = MlpModel(
input_size=input_dim,
hidden_sizes=hidden_sizes,
output_size=2 * delta_dim + 4, # 3 for each corners, times two for std, 4 probs
)
self.delta_distribution = Gaussian(
dim=delta_dim,
squash=True,
min_std=np.exp(MIN_LOG_STD),
max_std=np.exp(MAX_LOG_STD),
)
self.cat_distribution = Categorical(4)
def __init__(self,
input_classes: int,
output_size: int,
rnn_type: str = 'gru',
rnn_size: int = 256,
hidden_sizes: [List, Tuple] = None,
baselines_init: bool = True,
layer_norm: bool = False,
prev_action: int = 2,
prev_reward: int = 2,
):
super().__init__()
self._obs_dim = 0
self.rnn_is_lstm = rnn_type != 'gru'
self.preprocessor = tscr(OneHotLayer(input_classes))
rnn_class = get_rnn_class(rnn_type, layer_norm)
self.body = MlpModel(input_classes, hidden_sizes, None, nn.ReLU, None)
rnn_input_size = self.body.output_size
if prev_action: rnn_input_size += output_size # Use previous action as input
if prev_reward: rnn_input_size += 1 # Use previous reward as input
self.rnn = rnn_class(rnn_input_size, rnn_size) # Concat action, reward
self.pi = nn.Sequential(nn.ReLU(), nn.Linear(rnn_size, output_size), nn.Softmax(-1))
self.v = nn.Sequential(nn.ReLU(), nn.Linear(rnn_size, 1))
if baselines_init:
self.rnn.apply(apply_init); self.body.apply(apply_init)
self.pi.apply(partial(apply_init, gain=O_INIT_VALUES['pi']))
self.v.apply(partial(apply_init, gain=O_INIT_VALUES['v']))
self.body, self.pi, self.v = tscr(self.body), tscr(self.pi), tscr(self.v)
self.p_a = prev_action > 0
self.p_r = prev_reward > 0