I am working on a multi-agent reinforcement learning which each agent has its own q-network.However, the network cosists of several layers that share parameters among agents. I look up a lot of discussion about weights sharing and get confused now. I wonder if the code below can work? Thanks.
PS: the embedding layer and attention layer are the layers I want to share weights among agents.
def __init__(self, layer1, layer2, input_dim, output_dim): super().__init__() self.layer1 = layer1 self.layer2 = layer2 self.linear1 = nn.Linear(input_dim, output_dim) self.relu = nn.ReLU() def forward(self, x, neighbors): x_1 = self.layer1(x) neighbors_1 =  for neighbor in neighbors: neighbors_1.append(self.layer1(neighbor)) out = self.layer2(x_1, neighbors_1) out = self.linear1(out) out = self.relu(out) return out