def create(space, vocab_size):
  """Creates a SpaceSerializer for the given Gym space."""
  return {
      gym.spaces.Box: BoxSpaceSerializer,
      gym.spaces.Discrete: DiscreteSpaceSerializer,
      gym.spaces.MultiDiscrete: MultiDiscreteSpaceSerializer,
  }[type(space)](space, vocab_size)
Creates a SpaceSerializer for the given Gym space.
trax/rl/space_serializer.py