| 192 | return mu + eps * std |
| 193 | |
| 194 | def forward_conv(self, obs,flatten=True): |
| 195 | if obs.max() > 1.: |
| 196 | obs = obs / 255. |
| 197 | |
| 198 | time_step = obs.shape[1] // self.image_channel |
| 199 | obs = obs.view(obs.shape[0], time_step, self.image_channel, obs.shape[-2], obs.shape[-1]) |
| 200 | obs = obs.view(obs.shape[0]*time_step, self.image_channel, obs.shape[-2], obs.shape[-1]) |
| 201 | |
| 202 | self.outputs['obs'] = obs |
| 203 | conv = torch.relu(self.convs[0](obs)) |
| 204 | self.outputs['conv1'] = conv |
| 205 | |
| 206 | conv = torch.relu(self.convs[1](conv)) |
| 207 | self.outputs['conv%s' % (1 + 1)] = conv |
| 208 | |
| 209 | for i in range(2, self.num_layers): |
| 210 | conv = torch.relu(self.convs[i](conv)) |
| 211 | self.outputs['conv%s' % (i + 1)] = conv |
| 212 | |
| 213 | conv = conv.view(conv.size(0)//time_step, time_step, conv.size(1), conv.size(2), conv.size(3)) |
| 214 | |
| 215 | conv_current = conv[:, 1:, :, :, :] |
| 216 | conv_prev = conv_current - conv[:, :time_step-1, :, :, :].detach() |
| 217 | conv = torch.cat([conv_current, conv_prev], axis=1) |
| 218 | conv = conv.view(conv.size(0), conv.size(1)*conv.size(2), conv.size(3), conv.size(4)) |
| 219 | |
| 220 | if not flatten: |
| 221 | return conv |
| 222 | else: |
| 223 | conv = conv.view(conv.size(0), -1) |
| 224 | return conv |
| 225 | |
| 226 | |
| 227 | def forward(self, obs, detach=False): |