code: def policy(self, observations):
          
          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
docs: Calculate only action head for given state