How to implement this model in RLlib?

class cnn_model(nn.Module):
    def __init__(self, num_inputs, num_out, activation=nn.ReLU):
        super(cnn_model, self).__init__()
        self.conv1 = nn.Conv2d(num_inputs, 32, 3, stride=2, padding=1)
        self.conv2 = nn.Conv2d(32, 32, 3, stride=2, padding=1)
        self.conv3 = nn.Conv2d(32, 32, 3, stride=2, padding=1)
        self.conv4 = nn.Conv2d(32, 32, 3, stride=2, padding=1)

        self.fc1 = nn.Linear(32 * 6 * 6, 512)
        self.fc_out = nn.Linear(512, num_out)
        self._initialize_weights()

And the observation space of my env is (4,84,84), 4 is the stack size, so the num_inputs here is 4.

I haven’t used RLlib’s vision network myself, but I think it’s exactly what you want: RLlib Models, Preprocessors, and Action Distributions — Ray v1.2.0

There’s also a small example in the linked doc to show how you can configure the filters and layers of the ConvNet to whatever you need.

Other than that, here are more options for configuring the RLlib model: RLlib Models, Preprocessors, and Action Distributions — Ray v1.2.0

Does this help?

2 Likes