b9b54d4f30ff65cf1c54dc0cf90c938b48c44f90,reagent/models/dqn.py,DQNBase,forward,#DQNBase#Any#,41

Before Change

    def forward(self, input) -> rlt.AllActionQValues:
        q_values = self.dist(input).q_values
        if self.quantiles > 1:
            q_values = q_values.reshape((-1, self.action_dim, self.quantiles)).mean(
                dim=2
            )
        return rlt.AllActionQValues(q_values=q_values)

    def dist(self, input: rlt.PreprocessedState) -> rlt.AllActionQValues:
        if self.feature_extractor is not None:

After Change

    def input_prototype(self):
        return rlt.PreprocessedState.from_tensor(torch.randn(1, self.state_dim))

    def forward(self, input: rlt.PreprocessedState):
        q_values = self.fc(input.state.float_features)
        return rlt.AllActionQValues(q_values=q_values)


class _DistributedDataParallelFullyConnectedDQN(ModelBase):

Instances