def _inject_state(state_batch, state, env_idx):
    """
    Replace the state at the given index with a new state.
    """
    if state_batch is None:
        return
    elif isinstance(state_batch, tuple):
        return tuple(_inject_state(sb, s, env_idx)
                     for sb, s in zip(state_batch, state))
    state_batch[env_idx: env_idx+1] = state