(self, birdview, state)
| 119 | self.state_linear = nn.Sequential(*self.state_linear) |
| 120 | |
| 121 | def forward(self, birdview, state): |
| 122 | # birdview: [b, c, h, w] |
| 123 | # x = x.to(dtype=th.float32) / self.scale_ob |
| 124 | |
| 125 | for layer in self.stacks: |
| 126 | birdview = layer(birdview) |
| 127 | |
| 128 | x = th.flatten(birdview, 1) |
| 129 | x = th.relu(x) |
| 130 | |
| 131 | latent_state = self.state_linear(state) |
| 132 | |
| 133 | x = th.cat((x, latent_state), dim=1) |
| 134 | x = self.dense(x) |
| 135 | if self.final_relu: |
| 136 | x = th.relu(x) |
| 137 | return x |