(self, birdview, state)
| 96 | self.state_linear = nn.Sequential(*self.state_linear) |
| 97 | |
| 98 | def forward(self, birdview, state): |
| 99 | # birdview: [b, c, h, w] |
| 100 | # x = x.to(dtype=th.float32) / self.scale_ob |
| 101 | |
| 102 | for layer in self.stacks: |
| 103 | birdview = layer(birdview) |
| 104 | |
| 105 | x = th.flatten(birdview, 1) |
| 106 | x = th.relu(x) |
| 107 | |
| 108 | latent_state = self.state_linear(state) |
| 109 | |
| 110 | x = th.cat((x, latent_state), dim=1) |
| 111 | x = self.dense(x) |
| 112 | if self.final_relu: |
| 113 | x = th.relu(x) |
| 114 | return x |