(self, observation_space, chans=(16, 32, 32, 64, 64), states_neurons=[256],
features_dim=256, nblock=2, batch_norm=False, final_relu=True)
| 91 | |
| 92 | class ImpalaCNN(nn.Module): |
| 93 | def __init__(self, observation_space, chans=(16, 32, 32, 64, 64), states_neurons=[256], |
| 94 | features_dim=256, nblock=2, batch_norm=False, final_relu=True): |
| 95 | # (16, 32, 32) |
| 96 | super().__init__() |
| 97 | self.features_dim = features_dim |
| 98 | self.final_relu = final_relu |
| 99 | |
| 100 | # image encoder |
| 101 | curshape = observation_space['birdview'].shape |
| 102 | s = 1 / np.sqrt(len(chans)) # per stack scale |
| 103 | self.stacks = nn.ModuleList() |
| 104 | for outchan in chans: |
| 105 | stack = tu.CnnDownStack(curshape[0], nblock=nblock, outchan=outchan, scale=s, batch_norm=batch_norm) |
| 106 | self.stacks.append(stack) |
| 107 | curshape = stack.output_shape(curshape) |
| 108 | |
| 109 | # dense after concatenate |
| 110 | n_image_latent = tu.intprod(curshape) |
| 111 | self.dense = tu.NormedLinear(n_image_latent+states_neurons[-1], features_dim, scale=1.4) |
| 112 | |
| 113 | # state encoder |
| 114 | states_neurons = [observation_space['state'].shape[0]] + states_neurons |
| 115 | self.state_linear = [] |
| 116 | for i in range(len(states_neurons)-1): |
| 117 | self.state_linear.append(tu.NormedLinear(states_neurons[i], states_neurons[i+1])) |
| 118 | self.state_linear.append(nn.ReLU()) |
| 119 | self.state_linear = nn.Sequential(*self.state_linear) |
| 120 | |
| 121 | def forward(self, birdview, state): |
| 122 | # birdview: [b, c, h, w] |
nothing calls this directly
no test coverage detected