Choose the action space for the action proto.

def spatial(action, action_space):
  """Choose the action space for the action proto."""
  if action_space == ActionSpace.FEATURES:
    return action.action_feature_layer
  elif action_space == ActionSpace.RGB:
    return action.action_render
  else:
    raise ValueError("Unexpected value for action_space: %s" % action_space)