code: def spatial(action, action_space):
    
    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)
docs: Choose the action space for the action proto.