import torch import torch.nn as nn import torch.nn.functional as F def tie_weights(src, trg): assert type(src) == type(trg) trg.weight = src.weight trg.bias = src.bias OUT_DIM = {2: 39, 4: 35, 6: 31} class PixelEncoder(nn.Module): """Convolutional encoder of pixels observations.""" def __init__(self, obs_shape, feature_dim, num_layers=2, num_filters=32): super().__init__() assert len(obs_shape) == 3 self.feature_dim = feature_dim self.num_layers = num_layers self.convs = nn.ModuleList( [nn.Conv2d(obs_shape[0], num_filters, 3, stride=2)] ) for i in range(num_layers - 1): self.convs.append(nn.Conv2d(num_filters, num_filters, 3, stride=1)) out_dim = OUT_DIM[num_layers] self.fc = nn.Linear(num_filters * out_dim * out_dim, self.feature_dim * 2) self.ln = nn.LayerNorm(self.feature_dim * 2) self.outputs = dict() def reparameterize(self, mu, logstd): std = torch.exp(logstd) eps = torch.randn_like(std) return mu + eps * std def forward_conv(self, obs): obs = obs / 255. self.outputs['obs'] = obs conv = torch.relu(self.convs[0](obs)) self.outputs['conv1'] = conv for i in range(1, self.num_layers): conv = torch.relu(self.convs[i](conv)) self.outputs['conv%s' % (i + 1)] = conv h = conv.view(conv.size(0), -1) return h def forward(self, obs, detach=False): h = self.forward_conv(obs) if detach: h = h.detach() h_fc = self.fc(h) self.outputs['fc'] = h_fc h_norm = self.ln(h_fc) self.outputs['ln'] = h_norm h_tan = torch.tanh(h_norm) mu, logstd = torch.chunk(h_tan, 2, dim=-1) self.outputs['mu'] = mu self.outputs['logstd'] = logstd std = torch.tanh(h_norm) self.outputs['std'] = std out = self.reparameterize(mu, logstd) return out, mu, logstd def copy_conv_weights_from(self, source): """Tie convolutional layers""" # only tie conv layers for i in range(self.num_layers): tie_weights(src=source.convs[i], trg=self.convs[i]) def log(self, L, step, log_freq): if step % log_freq != 0: return for k, v in self.outputs.items(): L.log_histogram('train_encoder/%s_hist' % k, v, step) if len(v.shape) > 2: L.log_image('train_encoder/%s_img' % k, v[0], step) for i in range(self.num_layers): L.log_param('train_encoder/conv%s' % (i + 1), self.convs[i], step) L.log_param('train_encoder/fc', self.fc, step) L.log_param('train_encoder/ln', self.ln, step) class IdentityEncoder(nn.Module): def __init__(self, obs_shape, feature_dim, num_layers, num_filters): super().__init__() assert len(obs_shape) == 1 self.feature_dim = obs_shape[0] def forward(self, obs, detach=False): return obs def copy_conv_weights_from(self, source): pass def log(self, L, step, log_freq): pass class TransitionModel(nn.Module): def __init__(self, state_size, hidden_size, action_size, history_size): super().__init__() self.state_size = state_size self.hidden_size = hidden_size self.action_size = action_size self.history_size = history_size self.act_fn = nn.ELU() self.fc_state_action = nn.Linear(state_size + action_size, hidden_size) self.history_cell = nn.GRUCell(hidden_size, history_size) self.fc_state_mu = nn.Linear(history_size + hidden_size, state_size) self.fc_state_sigma = nn.Linear(history_size + hidden_size, state_size) self.batch_norm = nn.BatchNorm1d(hidden_size) self.batch_norm2 = nn.BatchNorm1d(state_size) self.min_sigma = 1e-4 self.max_sigma = 1e0 def init_states(self, batch_size, device): self.prev_state = torch.zeros(batch_size, self.state_size).to(device) self.prev_action = torch.zeros(batch_size, self.action_size).to(device) self.prev_history = torch.zeros(batch_size, self.history_size).to(device) def get_dist(self, mean, std): distribution = torch.distributions.Normal(mean, std) distribution = torch.distributions.independent.Independent(distribution, 1) return distribution def stack_states(self, states, dim=0): s = dict( mean = torch.stack([state['mean'] for state in states], dim=dim), std = torch.stack([state['std'] for state in states], dim=dim), sample = torch.stack([state['sample'] for state in states], dim=dim), history = torch.stack([state['history'] for state in states], dim=dim),) if 'distribution' in states: dist = dict(distribution = [state['distribution'] for state in states]) s.update(dist) return s def seq_to_batch(self, state, name): return dict( sample = torch.reshape(state[name], (state[name].shape[0]* state[name].shape[1], *state[name].shape[2:]))) def transition_step(self, prev_state, prev_action, prev_hist, prev_not_done): prev_state = prev_state.detach() * prev_not_done prev_hist = prev_hist * prev_not_done state_action_enc = self.fc_state_action(torch.cat([prev_state, prev_action], dim=-1)) state_action_enc = self.act_fn(self.batch_norm(state_action_enc)) current_hist = self.history_cell(state_action_enc, prev_hist) state_mu = self.act_fn(self.fc_state_mu(torch.cat([state_action_enc, prev_hist], dim=-1))) state_sigma = F.softplus(self.fc_state_sigma(torch.cat([state_action_enc, prev_hist], dim=-1))) sample_state = state_mu + torch.randn_like(state_mu) * state_sigma state_enc = {"mean": state_mu, "std": state_sigma, "sample": sample_state, "history": current_hist} return state_enc def observe_step(self, prev_state, prev_action, prev_history): state_action_enc = self.act_fn(self.batch_norm(self.fc_state_action(torch.cat([prev_state, prev_action], dim=-1)))) current_history = self.history_cell(state_action_enc, prev_history) state_mu = self.act_fn(self.batch_norm2(self.fc_state_mu(torch.cat([state_action_enc, prev_history], dim=-1)))) state_sigma = F.softplus(self.fc_state_sigma(torch.cat([state_action_enc, prev_history], dim=-1))) sample_state = state_mu + torch.randn_like(state_mu) * state_sigma state_enc = {"mean": state_mu, "std": state_sigma, "sample": sample_state, "history": current_history} return state_enc def observe_rollout(self, rollout_states, rollout_actions, init_history, nonterms): observed_rollout = [] for i in range(rollout_states.shape[0]): rollout_states_ = rollout_states[i] rollout_actions_ = rollout_actions[i] init_history_ = nonterms[i] * init_history state_enc = self.observe_step(rollout_states_, rollout_actions_, init_history_) init_history = state_enc["history"] observed_rollout.append(state_enc) observed_rollout = self.stack_states(observed_rollout, dim=0) return observed_rollout def reparemeterize(self, mean, std): eps = torch.randn_like(mean) return mean + eps * std def club_loss(x_samples, x_mu, x_logvar, y_samples): sample_size = x_samples.shape[0] random_index = torch.randperm(sample_size).long() positive = -(x_mu - y_samples)**2 / x_logvar.exp() negative = - (x_mu - y_samples[random_index])**2 / x_logvar.exp() upper_bound = (positive.sum(dim = -1) - negative.sum(dim = -1)).mean() return upper_bound/2. _AVAILABLE_ENCODERS = {'pixel': PixelEncoder, 'identity': IdentityEncoder} def make_encoder( encoder_type, obs_shape, feature_dim, num_layers, num_filters ): assert encoder_type in _AVAILABLE_ENCODERS return _AVAILABLE_ENCODERS[encoder_type]( obs_shape, feature_dim, num_layers, num_filters )