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