def map_args_obs_to_shape(obs, *args, **kwargs):
    return (unvectorize_shape_2d(obs), kwargs, {'obs': obs})
def map_args_obs_to_shape(obs, *args, **kwargs):
    return (unvectorize_shape_2d(obs), kwargs, {'obs': obs})