def _reduce_states(state_batch, env_idx):
    """
    Reduce a batch of states to a batch of one state.
    """
    if state_batch is None:
        return None
    elif isinstance(state_batch, tuple):
        return tuple(_reduce_states(s, env_idx) for s in state_batch)
    return state_batch[env_idx: env_idx+1].copy()