| 193 | |
| 194 | # Predict next hidden states given current states and actions |
| 195 | class DynamicsNetwork(nn.Module): |
| 196 | def __init__( |
| 197 | self, |
| 198 | num_blocks, |
| 199 | num_channels, |
| 200 | reduced_channels_reward, |
| 201 | fc_reward_layers, |
| 202 | full_support_size, |
| 203 | block_output_size_reward, |
| 204 | lstm_hidden_size=64, |
| 205 | momentum=0.1, |
| 206 | init_zero=False, |
| 207 | ): |
| 208 | """Dynamics network |
| 209 | Parameters |
| 210 | ---------- |
| 211 | num_blocks: int |
| 212 | number of res blocks |
| 213 | num_channels: int |
| 214 | channels of hidden states |
| 215 | fc_reward_layers: list |
| 216 | hidden layers of the reward prediction head (MLP head) |
| 217 | full_support_size: int |
| 218 | dim of reward output |
| 219 | block_output_size_reward: int |
| 220 | dim of flatten hidden states |
| 221 | lstm_hidden_size: int |
| 222 | dim of lstm hidden |
| 223 | init_zero: bool |
| 224 | True -> zero initialization for the last layer of reward mlp |
| 225 | """ |
| 226 | super().__init__() |
| 227 | self.num_channels = num_channels |
| 228 | self.lstm_hidden_size = lstm_hidden_size |
| 229 | |
| 230 | self.conv = conv3x3(num_channels, num_channels - 1) |
| 231 | self.bn = nn.BatchNorm2d(num_channels - 1, momentum=momentum) |
| 232 | self.resblocks = nn.ModuleList( |
| 233 | [ResidualBlock(num_channels - 1, num_channels - 1, momentum=momentum) for _ in range(num_blocks)] |
| 234 | ) |
| 235 | |
| 236 | self.reward_resblocks = nn.ModuleList( |
| 237 | [ResidualBlock(num_channels - 1, num_channels - 1, momentum=momentum) for _ in range(num_blocks)] |
| 238 | ) |
| 239 | |
| 240 | self.conv1x1_reward = nn.Conv2d(num_channels - 1, reduced_channels_reward, 1) |
| 241 | self.bn_reward = nn.BatchNorm2d(reduced_channels_reward, momentum=momentum) |
| 242 | self.block_output_size_reward = block_output_size_reward |
| 243 | self.lstm = nn.LSTM(input_size=self.block_output_size_reward, hidden_size=self.lstm_hidden_size) |
| 244 | self.bn_value_prefix = nn.BatchNorm1d(self.lstm_hidden_size, momentum=momentum) |
| 245 | self.fc = mlp(self.lstm_hidden_size, fc_reward_layers, full_support_size, init_zero=init_zero, momentum=momentum) |
| 246 | |
| 247 | def forward(self, x, reward_hidden): |
| 248 | state = x[:,:-1,:,:] |
| 249 | x = self.conv(x) |
| 250 | x = self.bn(x) |
| 251 | |
| 252 | x += state |