def encode_obs(self, obs):
obs_shape = obs.shape
per_image_shape = obs_shape[3:]
effetive_batch_size = get_product_of_iterable(obs_shape[:2])
num_frames = obs_shape[2]
h_t = self.convolutional_encoder(obs.view(-1, *per_image_shape)).view(
effetive_batch_size, num_frames, -1)
h_t = torch.mean(h_t, dim=1)
return h_t, effetive_batch_size
def get_weights_dict(self):
_latent_size = self.config.model.imagination_model.latent_size
_action_size = get_product_of_iterable(
self.config.env.action_space["shape"])
return torch.nn.ModuleDict({
"w_action":
torch.nn.Sequential(nn.Linear(_action_size, _latent_size)),
"w_z":
torch.nn.Sequential(nn.Linear(_latent_size, _latent_size)),
})
def __init__(self, config):
super().__init__(config=config)
self._input_size = self.config.model.imagination_model.hidden_state_size + self.config.model.imagination_model.latent_size
# Note that the hidden state size corresponds to the encoding of the observation in the pixel space.
self._output_size = get_product_of_iterable(
self.config.env.action_space.shape)
self.policy = nn.Sequential(
nn.Linear(self._input_size, self._input_size), nn.ReLU(),
nn.Linear(self._input_size, int(self._input_size / 2)), nn.ReLU(),
nn.Linear(int(self._input_size / 2), int(self._input_size / 4)),
nn.ReLU(), nn.Linear(int(self._input_size / 4), self._output_size))
self.criteria = nn.MSELoss()
def get_action_space_size(config):
return get_product_of_iterable(config.env.action_space.shape)
def loss(self, output, x):
""" loss function """
true_obs = x.next_obs
# Not that we have to manually divide because of an issue in Pytorch. The fix is available only in master for now.
return F.mse_loss(true_obs[:, :, 3, :, :, :], output) * 255 / (
get_product_of_iterable(output.shape))
def move_data_to_device(data, device):
if (type(data) in [list, tuple]):
return list(map(lambda x: x.to(device),
data)), get_product_of_iterable(data[1].shape[:2])
else:
return data.to(device), len(data)