Method __init__

utils/spatial_att_github.py:249–285

Agent implementing the attention agent.

(
        self,
        num_actions,
        hidden_size: int = 256,
        c_v: int = 120,
        c_k: int = 8,
        c_s: int = 64,
        num_queries: int = 4,
    )

class Agent(nn.Module):
    def __init__(
        self,
        num_actions,
        hidden_size: int = 256,
        c_v: int = 120,
        c_k: int = 8,
        c_s: int = 64,
        num_queries: int = 4,
    ):
        """Agent implementing the attention agent.
        """
        super(Agent, self).__init__()
        self.hidden_size = hidden_size
        self.c_v, self.c_k, self.c_s, self.num_queries = c_v, c_k, c_s, num_queries

        self.vision = VisionNetwork()
        self.query = QueryNetwork()
        # TODO: Implement SpatialBasis.
        self.spatial = SpatialBasis()

        self.answer_processor = nn.Sequential(
            # 1026 x 512
            nn.Linear(
                (c_v + c_s) * num_queries + (c_k + c_s) * num_queries + 1 + 1, 512
            ),
            nn.ReLU(),
            nn.Linear(512, hidden_size),
        )

        self.policy_core = nn.LSTMCell(hidden_size, hidden_size)


        self.prev_output = None
        self.prev_hidden = None

        self.policy_head = nn.Sequential(nn.Linear(hidden_size, num_actions))
        self.values_head = nn.Sequential(nn.Linear(hidden_size, num_actions))

    def reset(self):
        self.vision.reset()