956f94198f7eecbf9bc33704b9e7c0d39528bc52,catalyst/rl/agents/layers.py,StateActionNet,_forward_lama,#StateActionNet#Any#Any#,218

Before Change

        x = self.observation_net(x)

        x = x.view(batch_size, history_len, -1)
        observation_ = self.memory_net(x)

        // @TODO: add option to collapse observations based on action
        x = action.view(action.shape[0], -1)
        action_ = self.action_net(x)

        x = torch.cat((observation_, action_), dim=1)
        x = self.head_net(x)
        return x

After Change

        if len(obs_.shape) < 3:
            obs_ = obs_.unsqueeze(1)

        if self.observation_net is not None:
            batch_size, history_len, feature_size = obs_.shape
            obs_ = obs_.view(-1, feature_size)
            obs_ = self.observation_net(obs_)
            obs_ = obs_.view(batch_size, history_len, -1)
            obs_ = self.aggregation_net(obs_)

        // @TODO: add option to collapse observations based on action
        act_ = action.view(action.shape[0], -1)
        if self.action_net is not None:
            act_ = self.action_net(act_)

        x = torch.cat((obs_, act_), dim=1)
        x = self.main_net(x)
        x = self.head_net(x)
        return x

Instances