def __init__(self,
units: int = 64,
modalities: Set[str] = {"image", "pos", "sensors"}):
"""Initializes a measurement model for our door task."""
super().__init__(state_dim=2)
valid_modalities = {"image", "pos", "sensors"}
assert len(valid_modalities
| modalities) == 3, "Received invalid modality"
assert len(modalities) > 0, "Received empty modality list"
self.modalities = modalities
if "image" in modalities:
self.observation_image_layers = layers.observation_image_layers(
units, spanning_avg_pool=False)
if "pos" in modalities:
self.observation_pos_layers = layers.observation_pos_layers(units)
if "sensors" in modalities:
self.observation_sensors_layers = layers.observation_sensors_layers(
units)
self.state_layers = layers.state_layers(units)
self.shared_layers = nn.Sequential(
nn.Linear(units * (1 + len(modalities)), units),
nn.ReLU(inplace=True),
resblocks.Linear(units),
resblocks.Linear(units),
nn.Linear(units, 1),
# nn.LogSigmoid()
)
self.units = units
def __init__(self,
units: int = 64,
state_dim: int = 2,
know_image_blackout=False):
modality_count = 2
super().__init__(modality_count=modality_count, state_dim=state_dim)
weighting_types = ["sigmoid", "softmax", "absolute"]
self.observation_image_layers = layers.observation_image_layers(units)
self.observation_pos_layers = layers.observation_pos_layers(units)
self.observation_sensors_layers = layers.observation_sensors_layers(
units)
self.weighting_type = "sigmoid"
assert self.weighting_type in weighting_types
if self.weighting_type == "sigmoid":
# Initial fusion layers
self.fusion_layers = nn.Sequential(
nn.Linear(units * 3, units),
nn.ReLU(inplace=True),
resblocks.Linear(units),
nn.Linear(units, modality_count * self.state_dim),
nn.Sigmoid(),
)
else:
self.fusion_layers = nn.Sequential(
nn.Linear(units * 3, units),
nn.ReLU(inplace=True),
resblocks.Linear(units),
nn.Linear(units, modality_count * self.state_dim),
)
self.know_image_blackout = know_image_blackout
def make_feature_extractor(in_dim: int, out_dim: int,
hidden_units: int) -> nn.Module:
"""Helper to create a simple MLP model for embedding."""
return nn.Sequential(
nn.Linear(in_dim, hidden_units),
nn.LeakyReLU(inplace=True),
resblocks.Linear(hidden_units, activation="leaky_relu"),
resblocks.Linear(hidden_units, activation="leaky_relu"),
resblocks.Linear(hidden_units, activation="leaky_relu"),
nn.Linear(hidden_units, out_dim),
)
def __init__(self, state_dim=2, units=64, missing_modalities=[]):
super().__init__()
obs_pos_dim = 3
obs_sensors_dim = 7
self.state_dim = state_dim
# Missing modalities
self.modalities = set(["image", 'gripper_sensors', 'gripper_pos'])
if missing_modalities:
if type(missing_modalities) == list:
self.modalities -= set(missing_modalities)
else:
assert missing_modalities in self.modalities
self.modalities -= set([missing_modalities])
self.observation_image_layers = nn.Sequential(
nn.Conv2d(in_channels=1, out_channels=32, kernel_size=5,
padding=2),
nn.ReLU(inplace=True),
resblocks.Conv2d(channels=32, kernel_size=3),
nn.Conv2d(in_channels=32,
out_channels=16,
kernel_size=3,
padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(in_channels=16, out_channels=8, kernel_size=3,
padding=1),
nn.Flatten(), # 32 * 32 * 8
nn.Linear(8 * 32 * 32, units),
nn.ReLU(inplace=True),
resblocks.Linear(units),
)
self.observation_pos_layers = nn.Sequential(
nn.Linear(obs_pos_dim, units),
resblocks.Linear(units),
)
self.observation_sensors_layers = nn.Sequential(
nn.Linear(obs_sensors_dim, units),
resblocks.Linear(units),
)
self.state_layers = nn.Sequential(nn.Linear(state_dim, units), )
self.shared_layers = nn.Sequential(
nn.Linear(units * (len(self.modalities) + 1), units),
nn.ReLU(inplace=True),
resblocks.Linear(units),
resblocks.Linear(units),
nn.Linear(units, 1),
# nn.LogSigmoid()
)
self.units = units
def __init__(self,
state_dim=2,
state_noise_stddev=None,
units=32,
use_particles=True,
identity_prediction_dims=(),
learnable_Q=False):
super().__init__()
control_dim = 7
self.state_dim = state_dim
self.use_particles = use_particles
self.identity_prediction_dims = set(identity_prediction_dims)
assert len(self.identity_prediction_dims) == 0, "not implemented!"
if state_noise_stddev is not None:
self.state_noise_stddev = state_noise_stddev
else:
self.state_noise_stddev = self.default_state_noise_stddev
# Convert scalar noise values to arrays
if type(self.state_noise_stddev) == float:
self.state_noise_stddev = np.array([self.state_noise_stddev] *
state_dim)
self.state_layers = nn.Sequential(
nn.Linear(state_dim, units),
resblocks.Linear(units),
)
self.control_layers = nn.Sequential(
nn.Linear(control_dim, units),
resblocks.Linear(units),
)
self.shared_layers = nn.Sequential(
nn.Linear(units * 2, units),
resblocks.Linear(units),
resblocks.Linear(units),
resblocks.Linear(units),
# We add 1 to our desired output dimension to produce an extra
# "gate" term -- see implementation below
#
# We also subtract len(identity_prediction_dims) to avoid
# regressing updates for dimensions like mass, friction, etc
nn.Linear(units, (state_dim - len(identity_prediction_dims)) + 1),
)
self.units = units
Q_l = torch.from_numpy(np.array(self.state_noise_stddev)).float()
self.learnable_Q = learnable_Q
self.Q_l = torch.nn.Parameter(Q_l, requires_grad=self.learnable_Q)
def __init__(
self,
units: int = 64,
modalities: Set[str] = {"image", "pos", "sensors"},
add_R_noise: float = 1e-6,
noise_R_tril: torch.Tensor = None,
):
"""Initializes a measurement model for our door task."""
super().__init__(state_dim=3)
self.noise_R_tril = noise_R_tril
valid_modalities = {"image", "pos", "sensors"}
assert len(valid_modalities
| modalities) == 3, "Receivedi invalid modality"
assert len(modalities) > 0, "Received empty modality list"
self.modalities = modalities
if "image" in modalities:
self.observation_image_layers = layers.observation_image_layers(
units)
if "pos" in modalities:
self.observation_pos_layers = layers.observation_pos_layers(units)
if "sensors" in modalities:
self.observation_sensors_layers = layers.observation_sensors_layers(
units)
self.shared_layers = nn.Sequential(
nn.Linear(units * (len(modalities)), units * 2),
nn.ReLU(inplace=True),
resblocks.Linear(units * 2),
resblocks.Linear(units * 2),
# nn.Linear(units, units),
)
self.r_layer = nn.Sequential(
nn.Linear(units, self.state_dim),
nn.ReLU(inplace=True),
resblocks.Linear(self.state_dim),
nn.Linear(self.state_dim, self.state_dim),
)
self.z_layer = nn.Sequential(
nn.Linear(units, self.state_dim),
nn.ReLU(inplace=True),
resblocks.Linear(self.state_dim),
nn.Linear(self.state_dim, self.state_dim),
)
self.units = units
self.add_R_noise = torch.ones(self.state_dim) * add_R_noise
def __init__(self,
units=128,
state_dim=2,
use_states=False,
missing_modalities=None,
add_R_noise=1e-6):
super().__init__(units=units,
state_dim=state_dim,
use_states=False,
missing_modalities=missing_modalities,
add_R_noise=add_R_noise)
self.observation_image_layers = nn.Sequential(
nn.Conv2d(in_channels=1, out_channels=32, kernel_size=5,
padding=2),
nn.ReLU(inplace=True),
resblocks.Conv2d(channels=32, kernel_size=3),
nn.Conv2d(in_channels=32,
out_channels=16,
kernel_size=3,
padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(in_channels=16, out_channels=2, kernel_size=3,
padding=1),
)
#each avg pool gives us 16x2
self.gap_h = nn.AvgPool2d((32, 2))
self.gap_w = nn.AvgPool2d((2, 32))
self.gap_layers = nn.Sequential(
nn.Linear(32 * 2, units),
nn.ReLU(inplace=True),
resblocks.Linear(units),
)
def __init__(self,
units=128,
state_dim=2,
use_states=False,
missing_modalities=None,
add_R_noise=1e-6):
super().__init__(units, state_dim, False, missing_modalities,
add_R_noise)
self.observation_image_layers = nn.Sequential(
nn.Conv2d(in_channels=1, out_channels=32, kernel_size=5,
padding=2),
nn.ReLU(inplace=True),
resblocks.Conv2d(channels=32, kernel_size=3),
nn.Conv2d(in_channels=32,
out_channels=16,
kernel_size=3,
padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(in_channels=16, out_channels=2, kernel_size=3,
padding=1),
nn.AvgPool2d(5, 3),
nn.Flatten(),
nn.Linear(10 * 10 * 2, units),
nn.ReLU(inplace=True),
resblocks.Linear(units),
)
def __init__(self, units: int = 64):
"""Initializes an LSTM architecture for our door task."""
super().__init__(state_dim=3)
self.lstm_hidden_dim = 512
self.lstm_num_layers = 2
self.units = units
# Observation encoders
self.image_rows = 32
self.image_cols = 32
self.observation_image_layers = layers.observation_image_layers(units)
self.observation_pos_layers = layers.observation_pos_layers(units)
self.observation_sensors_layers = layers.observation_sensors_layers(
units)
# Control layers
self.control_layers = layers.control_layers(units)
# Initial fusion layers
self.fusion_layers = nn.Sequential(
nn.Linear(units * 4, units),
nn.ReLU(inplace=True),
resblocks.Linear(units),
)
# LSTM layers
self.lstm = nn.LSTM(units, self.lstm_hidden_dim, self.lstm_num_layers)
# Output layers
self.output_layers = nn.Sequential(
nn.Linear(self.lstm_hidden_dim, units),
nn.ReLU(inplace=True),
nn.Linear(units, self.state_dim),
)
def __init__(self, missing_modality, units=16, state_dim=2):
super().__init__(units, state_dim)
self.modalities = ["image", 'gripper_sensors', 'gripper_pos']
if type(missing_modality) == list:
self.modalities = [
mod for mod in self.modalities if mod not in missing_modality
]
else:
assert missing_modality in self.modalities
self.modalities.pop(self.modalities.index(missing_modality))
dims = (len(self.modalities))
self.shared_layers = nn.Sequential(
nn.Linear(units * dims, units * 2),
nn.ReLU(inplace=True),
resblocks.Linear(2 * units),
resblocks.Linear(2 * units),
)
def __init__(self, state_dim=2, state_noise_stddev=None):
super().__init__()
self.state_dim = state_dim
if state_noise_stddev is not None:
assert len(state_noise_stddev) == self.state_dim
self.state_noise_stddev = state_noise_stddev
else:
self.state_noise_stddev = np.array(
self.default_state_noise_stddev[:self.state_dim])
self.layer = nn.Sequential(
nn.Linear(self.state_dim, self.state_dim),
nn.ReLU(inplace=True),
resblocks.Linear(self.state_dim),
resblocks.Linear(self.state_dim),
)
self.Q = torch.from_numpy(np.diag(self.state_noise_stddev))
print("Currently deprecated. Use models in panda_models.py")
def __init__(self, know_image_blackout: bool, units: int = 64):
modality_count = 2
super().__init__(modality_count=modality_count)
self.know_image_blackout = know_image_blackout
self.observation_image_layers = layers.observation_image_layers(units)
self.observation_pos_layers = layers.observation_pos_layers(units)
self.observation_sensors_layers = layers.observation_sensors_layers(
units)
# Initial fusion layers
self.fusion_layers = nn.Sequential(
nn.Linear(units * 3, units),
nn.ReLU(inplace=True),
resblocks.Linear(units),
resblocks.Linear(units),
resblocks.Linear(units),
nn.Linear(units, modality_count),
# nn.LogSigmoid() # <- This drops performance slightly
)
def __init__(self, units=64):
"""Initializes a dynamics model for our door task."""
super().__init__(state_dim=2)
# Fixed dynamics covariance
self.Q_scale_tril = nn.Parameter(
torch.cholesky(torch.diag(torch.FloatTensor([0.02, 0.02]))),
requires_grad=False,
)
# Build neural network
self.state_layers = layers.state_layers(units=units)
self.control_layers = layers.control_layers(units=units)
self.shared_layers = nn.Sequential(
nn.Linear(units * 2, units),
resblocks.Linear(units),
resblocks.Linear(units),
resblocks.Linear(units),
nn.Linear(units, self.state_dim + 1),
)
self.units = units
def observation_sensors_layers(units: int) -> nn.Module:
"""Create an F/T sensor encoder block.
Args:
units (int): # of hidden units in network layers.
Returns:
nn.Module: Encoder block.
"""
return nn.Sequential(
nn.Linear(obs_sensors_dim, units),
nn.ReLU(inplace=True),
resblocks.Linear(units),
)
def control_layers(units: int) -> nn.Module:
"""Create a control command encoder block.
Args:
units (int): # of hidden units in network layers.
Returns:
nn.Module: Encoder block.
"""
return nn.Sequential(
nn.Linear(control_dim, units),
nn.ReLU(inplace=True),
resblocks.Linear(units),
)
def observation_image_layers(units: int, spanning_avg_pool: bool = False) -> nn.Module:
"""Create an image encoder block.
Args:
units (int): # of hidden units in network layers.
Returns:
nn.Module: Encoder block.
"""
if spanning_avg_pool:
# Architecture with full width/height average pools
return nn.Sequential(
nn.Conv2d(in_channels=1, out_channels=32, kernel_size=5, padding=2),
nn.ReLU(inplace=True),
resblocks.Conv2d(channels=32, kernel_size=3),
nn.Conv2d(in_channels=32, out_channels=16, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(in_channels=16, out_channels=2, kernel_size=3, padding=1),
_DualSpanningAvgPool(rows=32, cols=32, reduce_size=2),
nn.Linear(32 * 2, units),
nn.ReLU(inplace=True),
resblocks.Linear(units),
)
else:
# Default model
return nn.Sequential(
nn.Conv2d(in_channels=1, out_channels=32, kernel_size=5, padding=2),
nn.ReLU(inplace=True),
resblocks.Conv2d(channels=32, kernel_size=3),
nn.Conv2d(in_channels=32, out_channels=16, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(in_channels=16, out_channels=8, kernel_size=3, padding=1),
nn.Flatten(), # 32 * 32 * 8
nn.Linear(8 * 32 * 32, units),
nn.ReLU(inplace=True),
resblocks.Linear(units),
)
def test_resblock_smoke_test():
"""Make sure we can build all resblocks."""
for inplace in (True, False):
for activation in resblocks.Base._activation_types.keys():
resblocks.Linear(20,
3,
activation=activation,
activations_inplace=inplace)
resblocks.Conv2d(
channels=20,
bottleneck_channels=3,
kernel_size=5,
activation=activation,
activations_inplace=inplace,
)
def __init__(self, use_prev_state=True, units=32):
super().__init__()
self.use_prev_state = use_prev_state
self.units = units
obs_pos_dim = 7
obs_sensors_dim = 7
state_dim = 2
control_dim = 7
self.state_layers = nn.Sequential(
nn.Linear(state_dim, units // 2),
nn.ReLU(inplace=True),
resblocks.Linear(units // 2),
)
self.control_layers = nn.Sequential(
nn.Linear(control_dim, units // 2),
nn.ReLU(inplace=True),
resblocks.Linear(units // 2),
)
self.observation_image_layers = nn.Sequential(
nn.Conv2d(in_channels=1, out_channels=4, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
resblocks.Conv2d(channels=4),
nn.Conv2d(in_channels=4, out_channels=1, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Flatten(), # 32 * 32 = 1024
nn.Linear(1024, units),
nn.ReLU(inplace=True),
resblocks.Linear(units),
)
self.observation_pose_layers = nn.Sequential(
nn.Linear(obs_pos_dim, units),
resblocks.Linear(units),
)
self.observation_sensors_layers = nn.Sequential(
nn.Linear(obs_sensors_dim, units),
resblocks.Linear(units),
)
self.shared_layers = nn.Sequential(
nn.Linear((units // 2) * 2 + units * 3, units),
nn.ReLU(inplace=True),
resblocks.Linear(units),
resblocks.Linear(units),
nn.Linear(units, state_dim), # Directly output new state
# nn.LogSigmoid()
)
def observation_image_layers(units: int) -> nn.Module:
"""Create an image encoder block.
Args:
units (int): # of hidden units in network layers.
Returns:
nn.Module: Encoder block.
"""
return nn.Sequential(
nn.Conv2d(in_channels=1, out_channels=32, kernel_size=5, padding=2),
nn.ReLU(inplace=True),
resblocks.Conv2d(channels=32, kernel_size=3),
nn.Conv2d(in_channels=32, out_channels=16, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(in_channels=16, out_channels=8, kernel_size=3, padding=1),
nn.Flatten(), # 32 * 32 * 8
nn.Linear(8 * 32 * 32, units),
nn.ReLU(inplace=True),
resblocks.Linear(units),
)
def __init__(self,
state_dim=2,
units=32,
use_softmax=True,
use_log_softmax=False):
super().__init__()
obs_pose_dim = 3
obs_sensors_dim = 7
self.state_dim = state_dim
self.use_softmax = use_softmax
self.use_log_softmax = use_log_softmax
self.observation_image_layers = nn.Sequential(
nn.Conv2d(in_channels=1, out_channels=32, kernel_size=5,
padding=2),
nn.ReLU(inplace=True),
resblocks.Conv2d(channels=32, kernel_size=3),
nn.Conv2d(in_channels=32,
out_channels=16,
kernel_size=3,
padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(in_channels=16, out_channels=2, kernel_size=3,
padding=1),
nn.Flatten(), # 32 * 32 * 8
nn.Linear(2 * 32 * 32, units),
nn.ReLU(inplace=True),
resblocks.Linear(units),
)
self.observation_pose_layers = nn.Sequential(
nn.Linear(obs_pose_dim, units),
resblocks.Linear(units, activation='leaky_relu'),
)
self.observation_sensors_layers = nn.Sequential(
nn.Linear(obs_sensors_dim, units),
resblocks.Linear(units, activation='leaky_relu'),
)
# todo: the +1 only works for state dim =2
# it should be self.state_dim + (state_dim)(state_dim-1)/2
if self.use_softmax:
self.shared_layers = nn.Sequential(
nn.Linear(units * 3, units),
nn.ReLU(inplace=True),
resblocks.Linear(units, units),
resblocks.Linear(units, units),
resblocks.Linear(units, units),
nn.Linear(units, 2 * (self.state_dim + 1)),
)
else:
self.shared_layers = nn.Sequential(
nn.Linear(units * 3, units * 3),
nn.ReLU(inplace=True),
resblocks.Linear(3 * units),
)
self.force_prop_layer = nn.Sequential(
nn.Linear(units, units),
nn.ReLU(inplace=True),
nn.Linear(units, self.state_dim),
nn.Sigmoid(),
)
self.image_prop_layer = nn.Sequential(
nn.Linear(units, units),
nn.ReLU(inplace=True),
nn.Linear(units, self.state_dim),
nn.Sigmoid(),
)
for m in self.modules():
if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
nn.init.kaiming_normal_(m.weight.data)
if m.bias is not None:
m.bias.data.zero_()
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
self.units = units
def __init__(self,
units=128,
state_dim=2,
use_states=False,
missing_modalities=None,
add_R_noise=1e-6):
super().__init__()
obs_pose_dim = 3
obs_sensors_dim = 7
image_dim = (32, 32)
#We do not use states for EKF
self.state_dim = state_dim
# if we want to use states in measurement model update
self.use_states = use_states
# Missing modalities
self.modalities = set(["image", 'gripper_sensors', 'gripper_pos'])
if missing_modalities:
if type(missing_modalities) == list:
self.modalities -= set(missing_modalities)
else:
assert missing_modalities in self.modalities
self.modalities -= set([missing_modalities])
self.observation_image_layers = nn.Sequential(
nn.Conv2d(in_channels=1, out_channels=32, kernel_size=5,
padding=2),
nn.ReLU(inplace=True),
resblocks.Conv2d(channels=32, kernel_size=3),
nn.Conv2d(in_channels=32,
out_channels=16,
kernel_size=3,
padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(in_channels=16, out_channels=2, kernel_size=3,
padding=1),
nn.Flatten(), # 32 * 32 * 2
nn.Linear(2 * 32 * 32, units),
nn.ReLU(inplace=True),
resblocks.Linear(units),
)
self.observation_pose_layers = nn.Sequential(
nn.Linear(obs_pose_dim, units),
# nn.ReLU(inplace=True),
resblocks.Linear(units),
)
self.observation_sensors_layers = nn.Sequential(
nn.Linear(obs_sensors_dim, units),
# nn.ReLU(inplace=True),
resblocks.Linear(units),
)
# missing modalities
self.shared_layers = nn.Sequential(
nn.Linear(units * (len(self.modalities)), units * 2),
nn.ReLU(inplace=True),
resblocks.Linear(2 * units),
resblocks.Linear(2 * units),
)
self.r_layer = nn.Sequential(
nn.Linear(units, self.state_dim),
nn.ReLU(inplace=True),
resblocks.Linear(self.state_dim),
nn.Linear(self.state_dim, self.state_dim),
)
self.z_layer = nn.Sequential(
nn.Linear(units, self.state_dim),
nn.ReLU(inplace=True),
resblocks.Linear(self.state_dim),
nn.Linear(self.state_dim, self.state_dim),
)
self.units = units
# for m in self.modules():
# if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
# nn.init.kaiming_normal_(m.weight.data)
# if m.bias is not None:
# m.bias.data.zero_()
# elif isinstance(m, nn.BatchNorm2d):
# m.weight.data.fill_(1)
# m.bias.data.zero_()
self.add_R_noise = torch.ones(state_dim) * add_R_noise
def __init__(self, units=32):
obs_pos_dim = 3
obs_sensors_dim = 7
control_dim = 7
self.state_dim = 2
super().__init__()
self.lstm_hidden_dim = 4
self.lstm_num_layers = 2
self.units = units
# Observation encoders
self.image_rows = 32
self.image_cols = 32
self.observation_image_layers = nn.Sequential(
nn.Conv2d(in_channels=1, out_channels=32, kernel_size=5,
padding=2),
nn.ReLU(inplace=True),
resblocks.Conv2d(channels=32, kernel_size=3),
nn.Conv2d(in_channels=32,
out_channels=16,
kernel_size=3,
padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(in_channels=16, out_channels=8, kernel_size=3,
padding=1),
nn.Flatten(), # 32 * 32 * 8
nn.Linear(8 * 32 * 32, units),
nn.ReLU(inplace=True),
resblocks.Linear(units),
)
self.observation_pose_layers = nn.Sequential(
nn.Linear(obs_pos_dim, units),
nn.ReLU(inplace=True),
resblocks.Linear(units),
)
self.observation_sensors_layers = nn.Sequential(
nn.Linear(obs_sensors_dim, units),
nn.ReLU(inplace=True),
resblocks.Linear(units),
)
# Control layers
self.control_layers = nn.Sequential(
nn.Linear(control_dim, units),
nn.ReLU(inplace=True),
resblocks.Linear(units),
)
# Fusion layer
self.fusion_layers = nn.Sequential(
nn.Linear(units * 4, units),
nn.ReLU(inplace=True),
resblocks.Linear(units),
)
# LSTM layers
self.lstm = nn.LSTM(units,
self.lstm_hidden_dim,
self.lstm_num_layers,
batch_first=True)
# Define the output layer
self.output_layers = nn.Sequential(
nn.Linear(self.lstm_hidden_dim, units),
nn.ReLU(inplace=True),
# resblocks.Linear(units),
nn.Linear(units, self.state_dim),
)
def __init__(self,
units=16,
state_dim=2,
use_states=False,
use_spatial_softmax=False):
super().__init__()
print("Currently deprecated. Use models in panda_models.py")
obs_pose_dim = 3
obs_sensors_dim = 7
image_dim = (32, 32)
self.state_dim = state_dim
self.use_states = use_states
if self.use_states:
shared_layer_dim = units * 4
else:
shared_layer_dim = units * 3
if use_spatial_softmax:
self.observation_image_layers = nn.Sequential(
nn.Conv2d(in_channels=1,
out_channels=32,
kernel_size=5,
padding=2),
nn.ReLU(inplace=True),
resblocks.Conv2d(channels=32, kernel_size=3),
nn.Conv2d(in_channels=32,
out_channels=16,
kernel_size=3,
padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(in_channels=16,
out_channels=16,
kernel_size=3,
padding=1),
spatial_softmax.SpatialSoftmax(32, 32, 16),
nn.Linear(16 * 2, units),
nn.ReLU(inplace=True),
resblocks.Linear(units),
)
else:
self.observation_image_layers = nn.Sequential(
nn.Conv2d(in_channels=1,
out_channels=32,
kernel_size=5,
padding=2),
nn.ReLU(inplace=True),
resblocks.Conv2d(channels=32, kernel_size=3),
nn.Conv2d(in_channels=32,
out_channels=16,
kernel_size=3,
padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(in_channels=16,
out_channels=8,
kernel_size=3,
padding=1),
nn.Flatten(), # 32 * 32 = 1024
nn.Linear(8 * 32 * 32, units),
nn.ReLU(inplace=True),
resblocks.Linear(units),
)
self.observation_pose_layers = nn.Sequential(
nn.Linear(obs_pose_dim, units),
resblocks.Linear(units, activation='leaky_relu'),
)
self.observation_sensors_layers = nn.Sequential(
nn.Linear(obs_sensors_dim, units),
resblocks.Linear(units, activation='leaky_relu'),
)
self.state_layers = nn.Sequential(nn.Linear(self.state_dim, units), )
self.shared_layers = nn.Sequential(
nn.Linear(shared_layer_dim, units * 2),
nn.ReLU(inplace=True),
resblocks.Linear(2 * units),
resblocks.Linear(2 * units),
)
self.r_layer = nn.Sequential(
nn.Linear(units, self.state_dim),
nn.ReLU(inplace=True),
resblocks.Linear(self.state_dim),
)
self.z_layer = nn.Sequential(
nn.Linear(units, self.state_dim),
nn.ReLU(inplace=True),
resblocks.Linear(self.state_dim),
)
self.units = units
