(self, observation_space, chans=(16, 32, 32, 64, 64), states_neurons=[256],
features_dim=256, nblock=2, batch_norm=False, final_relu=True)
| 68 | |
| 69 | class ImpalaCNN(nn.Module): |
| 70 | def __init__(self, observation_space, chans=(16, 32, 32, 64, 64), states_neurons=[256], |
| 71 | features_dim=256, nblock=2, batch_norm=False, final_relu=True): |
| 72 | # (16, 32, 32) |
| 73 | super().__init__() |
| 74 | self.features_dim = features_dim |
| 75 | self.final_relu = final_relu |
| 76 | |
| 77 | # image encoder |
| 78 | curshape = observation_space['birdview'].shape |
| 79 | s = 1 / np.sqrt(len(chans)) # per stack scale |
| 80 | self.stacks = nn.ModuleList() |
| 81 | for outchan in chans: |
| 82 | stack = tu.CnnDownStack(curshape[0], nblock=nblock, outchan=outchan, scale=s, batch_norm=batch_norm) |
| 83 | self.stacks.append(stack) |
| 84 | curshape = stack.output_shape(curshape) |
| 85 | |
| 86 | # dense after concatenate |
| 87 | n_image_latent = tu.intprod(curshape) |
| 88 | self.dense = tu.NormedLinear(n_image_latent+states_neurons[-1], features_dim, scale=1.4) |
| 89 | |
| 90 | # state encoder |
| 91 | states_neurons = [observation_space['state'].shape[0]] + states_neurons |
| 92 | self.state_linear = [] |
| 93 | for i in range(len(states_neurons)-1): |
| 94 | self.state_linear.append(tu.NormedLinear(states_neurons[i], states_neurons[i+1])) |
| 95 | self.state_linear.append(nn.ReLU()) |
| 96 | self.state_linear = nn.Sequential(*self.state_linear) |
| 97 | |
| 98 | def forward(self, birdview, state): |
| 99 | # birdview: [b, c, h, w] |
nothing calls this directly
no test coverage detected