def __init__(
self,
recurrent,
embedding_buckets,
hidden_size=256,
img_mean=[0, 0, 0],
img_std=[1, 1, 1],
input_shape=(90, 160),
):
super(CNNBase, self).__init__(recurrent, hidden_size, hidden_size)
init_ = lambda m: init(
m,
nn.init.orthogonal_,
lambda x: nn.init.constant_(x, 0),
nn.init.calculate_gain("relu"),
)
self.img_mean = img_mean
self.img_std = img_std
imH, imW = input_shape
embedding_size = 0
# Assumes input size is (imH, imW)
self.rgb_encoder = nn.Sequential(
init_(nn.Conv2d(3, 2, 3, padding=1)),
nn.BatchNorm2d(2),
nn.ReLU(),
nn.MaxPool2d(2, stride=2),
init_(nn.Conv2d(2, 1, 3, padding=1)),
nn.BatchNorm2d(1),
nn.ReLU(),
nn.MaxPool2d(2, stride=2),
init_(nn.Conv2d(1, 1, 3, padding=1)),
nn.BatchNorm2d(1),
nn.ReLU(),
nn.MaxPool2d(2, stride=2),
Flatten(),
)
rgb_encoder_output = self.rgb_encoder(torch.randn(1, 3, imH, imW))
embedding_size += rgb_encoder_output.shape[1]
embedding_size += self._create_embeddings(embedding_buckets)
init_ = lambda m: init(m, nn.init.orthogonal_, lambda x: nn.init.
constant_(x, 0))
self.fuse_embedding = init_(nn.Linear(embedding_size, hidden_size))
self.critic_linear = init_(nn.Linear(hidden_size, 1))
self.train()
def __init__(self, config):
super().__init__()
self.nactions = config.nactions
self.hidden_size = config.hidden_size
embedding_buckets = config.EMBEDDING_BUCKETS
self.base = CNNBase(
True,
embedding_buckets,
hidden_size=self.hidden_size,
img_mean=config.NORMALIZATION.img_mean,
img_std=config.NORMALIZATION.img_std,
input_shape=config.image_scale_hw,
)
self.dist = Categorical(self.hidden_size, self.nactions)
init_ = lambda m: init(m, nn.init.orthogonal_, lambda x: nn.init.
constant_(x, 0))