def reset_weights(self):
        """ Initialize properly model weights """
        self.input_block.reset_weights()

        self.policy_backbone.reset_weights()
        self.value_backbone.reset_weights()

        self.action_head.reset_weights()
        self.value_head.reset_weights()