code: def extract_action_values(self, prediction):
      return prediction.squeeze()
  
file_path: coach/rl_coach/agents/imitation_agent.py
imports: from collections import OrderedDict
  from typing import Union
  from rl_coach.agents.agent import Agent
  from rl_coach.core_types import RunPhase, ActionInfo
  from rl_coach.logger import screen
  from rl_coach.spaces import DiscreteActionSpace
label_desc: Use image node_id 3 for calling the ImitationAgent obj's underlying member method code with example usage: obj.extract_action_values(prediction) and returns: prediction