def __init__()

in agents/obj_nets.py [0:0]

• 121 lines of code
• 3 McCabe index (conditional complexity)

    def __init__(self,
                 n_hist_frames=3,
                 n_fwd_times=0,
                 n_heads=8,
                 n_layers=6,
                 embedding_size=128,
                 aggregate_hidden_size=128,
                 score_hidden_first_layer_scale=1,
                 aggregate='mean',
                 n_inp=phyre_simulator.MAX_NUM_OBJECTS,
                 embed_tf=True,
                 shuffle_embed=False,
                 class_extra_layer=False,
                 tf_layer_norm=False):
        super().__init__()
        self.n_inp = n_inp
        self.score_hidden_first_layer_scale = score_hidden_first_layer_scale
        self.n_timesteps = n_hist_frames + n_fwd_times
        self.n_heads = n_heads
        self.n_layers = n_layers
        self.embedding_size = embedding_size
        self.embed_tf = embed_tf
        self.shuffle_embed = shuffle_embed
        self.tf_layer_norm = tf_layer_norm
        self.encoder = nn.Sequential(
            nn.Linear(phyre.FeaturizedObjects._NUM_FEATURES,
                      self.embedding_size),
            nn.ReLU(),
            nn.Linear(self.embedding_size, self.embedding_size),
        )
        self.pos_encoder = PositionalEncoding(self.embedding_size,
                                              max_len=self.n_timesteps,
                                              num_inp=self.n_inp)
        encoder_layers = TransformerEncoderLayer(self.embedding_size,
                                                 nhead=self.n_heads)
        if self.tf_layer_norm:
            norm = nn.LayerNorm(self.embedding_size)
        else:
            norm = None
        self.transformer_encoder = TransformerEncoder(encoder_layers,
                                                      num_layers=self.n_layers,
                                                      norm=norm)
        self.aggregate_hidden_size = aggregate_hidden_size

        self.aggregation = aggregate
        self.class_extra_layer = class_extra_layer
        self.score_obj = nn.Sequential(
            nn.Linear(self.embedding_size, self.aggregate_hidden_size),
            nn.ReLU(),
            nn.Linear(self.aggregate_hidden_size, self.aggregate_hidden_size),
            nn.ReLU(),
            nn.Linear(self.aggregate_hidden_size, 1),
        )
        self.score = nn.Sequential(
            nn.Linear(
                phyre_simulator.MAX_NUM_OBJECTS * self.embedding_size *
                self.n_timesteps, self.aggregate_hidden_size),
            nn.ReLU(),
            nn.Linear(self.aggregate_hidden_size, self.aggregate_hidden_size),
            nn.ReLU(),
            nn.Linear(self.aggregate_hidden_size, 1),
        )
        if self.class_extra_layer:
            self.score_goal = nn.Sequential(
                nn.Linear(
                    2 * self.embedding_size * self.n_timesteps,
                    int(self.embedding_size * self.n_timesteps *
                        phyre_simulator.MAX_NUM_OBJECTS / 2.0)),
                nn.ReLU(),
                nn.Linear(
                    int(self.embedding_size * self.n_timesteps *
                        phyre_simulator.MAX_NUM_OBJECTS / 2.0),
                    self.aggregate_hidden_size),
                nn.ReLU(),
                nn.Linear(self.aggregate_hidden_size,
                          self.aggregate_hidden_size),
                nn.ReLU(),
                nn.Linear(self.aggregate_hidden_size, 1),
            )
            self.score_timestep_embeddings = nn.Sequential(
                nn.Linear(
                    self.embedding_size * self.n_timesteps,
                    self.embedding_size * self.n_timesteps *
                    phyre_simulator.MAX_NUM_OBJECTS),
                nn.ReLU(),
                nn.Linear(
                    self.embedding_size * self.n_timesteps *
                    phyre_simulator.MAX_NUM_OBJECTS,
                    self.aggregate_hidden_size),
                nn.ReLU(),
                nn.Linear(self.aggregate_hidden_size,
                          self.aggregate_hidden_size),
                nn.ReLU(),
                nn.Linear(self.aggregate_hidden_size, 1),
            )
        else:
            self.score_goal = nn.Sequential(
                nn.Linear(
                    2 * self.embedding_size * self.n_timesteps,
                    self.aggregate_hidden_size *
                    self.score_hidden_first_layer_scale),
                nn.ReLU(),
                nn.Linear(
                    self.aggregate_hidden_size *
                    self.score_hidden_first_layer_scale,
                    self.aggregate_hidden_size),
                nn.ReLU(),
                nn.Linear(self.aggregate_hidden_size, 1),
            )
            self.score_timestep_embeddings = nn.Sequential(
                nn.Linear(
                    self.embedding_size * self.n_timesteps,
                    self.aggregate_hidden_size *
                    self.score_hidden_first_layer_scale),
                nn.ReLU(),
                nn.Linear(
                    self.aggregate_hidden_size *
                    self.score_hidden_first_layer_scale,
                    self.aggregate_hidden_size),
                nn.ReLU(),
                nn.Linear(self.aggregate_hidden_size, 1),
            )