Calculate only action head for given state

def policy(self, observations):
        """ Calculate only action head for given state """
        input_data = self.input_block(observations)
        policy_base_output = self.policy_backbone(input_data)
        policy_params = self.action_head(policy_base_output)
        return policy_params