code: def _optimal_deep_sea_policy(
      environment: deep_sea.DeepSea, timestep: dm_env.TimeStep
  ):
      action = environment._action_mapping[
          np.where(timestep.observation)
      ]  # pylint: disable=protected-access
      return action[0].astype(np.int32)
  
imports: from typing import Any, List
  from absl import flags
  from bsuite.environments import deep_sea
  import dm_env
  import numpy
  import tensorflow
  import tree
label_desc: Use image node_id 2 for calling a global function with example usage: _optimal_deep_sea_policy(environment, timestep) and returns: action