Vectorized Envs¶

def reset_done(
    self, 
    dones: Union[List[bool], np.ndarray],
    *, 
    seed: int | None = None, 
    options: Dict[str, Any] | None = None
):
    """
    Reset only the environments indicated by ``dones``.

    Args:
        dones: Boolean mask/list of length :attr:`n_envs`. Entries set to
            ``True`` are reset.
        seed: Optional base seed. Implementations may offset by environment index.
        options: Reset options forwarded to underlying environments.

    Returns:
        A tuple ``(reset_obs, reset_infos)`` where:

        * ``reset_obs`` is a list of length :attr:`n_envs` containing reset
          observations at indices that were reset, and ``None`` elsewhere.
        * ``reset_infos`` is a list of length :attr:`n_envs` containing reset
          info dicts at indices that were reset, and empty dicts elsewhere.

    Raises:
        NotImplementedError: If not implemented by a subclass.
    """
    raise NotImplementedError

def reset(self, *, seed: int | None = None, options: Dict[str, Any] | None = None):
    """
    Reset and return vectorized lists of length 1.

    Args:
        seed: Random seed forwarded to the underlying environment.
        options: Reset options forwarded to the underlying environment.

    Returns:
        ``([obs], [info])``.
    """
    obs, info = self.env.reset(seed=seed, options=options)
    return [obs], [info]

def reset_done(
    self, 
    dones: Union[List[bool], np.ndarray],
    *, 
    seed: int | None = None, 
    options: Dict[str, Any] | None = None
):
    """
    Conditionally reset the single environment.

    Args:
        dones: A length-1 mask. If ``dones[0]`` is ``True``, reset.
        seed: Random seed forwarded to the underlying environment.
        options: Reset options forwarded to the underlying environment.

    Returns:
        A pair ``(reset_obs, reset_infos)`` as described by
        :meth:`VecEnvWrapperBase.reset_done`.
    """
    dones = list(dones)
    assert len(dones) == 1
    if dones[0]:
        return self.reset(seed=seed, options=options)
    else:
        [None], [{}]

def step(self, action):
    """
    Step and return vectorized lists of length 1.

    Args:
        action: Action for the single environment.

    Returns:
        ``([obs], [rew], [terminated], [truncated], [info])``.
    """
    obs, rew, term, trunc, info = self.env.step(action)
    return [obs], [rew], [term], [trunc], [info]

def reset(self, *, seed: int | None = None, options: Dict[str, Any] | None = None):
    """
    Reset all environments and return lists.

    Args:
        seed: Optional base seed. If provided, environment ``i`` receives ``seed + i``.
        options: Reset options forwarded to each environment.

    Returns:
        A pair ``(obs_list, info_list)`` of length :attr:`n_envs`.
    """
    obs_list, info_list = [], []
    for i, env in enumerate(self.envs):
        s = None if seed is None else seed + i
        obs, info = env.reset(seed=s, options=options)
        obs_list.append(obs)
        info_list.append(info)
    return obs_list, info_list

def reset_done(
    self, 
    dones: Union[List[bool], np.ndarray],
    *, 
    seed: int | None = None, 
    options: Dict[str, Any] | None = None
):
    """
    Reset only environments whose done flag is True.

    Args:
        dones: Boolean mask/list of length :attr:`n_envs`.
        seed: Optional base seed. If provided, environment ``i`` receives ``seed + i``.
        options: Reset options forwarded to environments being reset.

    Returns:
        A tuple ``(reset_obs, reset_infos)`` where non-reset indices contain
        ``None`` and ``{}`` respectively.
    """
    dones = list(dones)
    assert len(dones) == self.n_envs

    reset_obs = [None] * self.n_envs
    reset_infos = [{} for _ in range(self.n_envs)]

    for i, done in enumerate(dones):
        if done:
            s = None if seed is None else seed + i
            obs, info = self.envs[i].reset(seed=s, options=options)
            reset_obs[i] = obs
            reset_infos[i] = info

    return reset_obs, reset_infos

def step(self, actions):
    """
    Step all environments.

    Args:
        actions: Iterable of actions of length :attr:`n_envs`.

    Returns:
        A 5-tuple of lists ``(obs_list, rew_list, term_list, trunc_list, info_list)``.

    Notes:
        The loop expects one action per environment. If the provided
        ``actions`` length mismatches :attr:`n_envs`, Python will raise.
    """
    obs_list, rew_list, term_list, trunc_list, info_list = [], [], [], [], []
    for env, action in zip(self.envs, actions):
        obs, rew, term, trunc, info = env.step(action)
        obs_list.append(obs)
        rew_list.append(rew)
        term_list.append(term)
        trunc_list.append(trunc)
        info_list.append(info)

    return obs_list, rew_list, term_list, trunc_list, info_list

def _normalize_obs(self, obs_list: List[np.ndarray]) -> List[np.ndarray]:
    """
    Normalize and clip a list of observations.

    Args:
        obs_list: List of raw observations of length :attr:`n_envs`.

    Returns:
        List of normalized observations.
    """
    obs_arr = np.asarray(obs_list, dtype=np.float32)
    norm = (obs_arr - self.obs_rms.mean) / np.sqrt(self.obs_rms.var + self.eps)
    norm = np.clip(norm, -self.clip_obs, self.clip_obs)
    return norm.tolist()

def _normalize_rew(self, rew_list: List[float]) -> List[float]:
    """
    Normalize and clip a list of rewards.

    Args:
        rew_list: List of raw rewards of length :attr:`n_envs`.

    Returns:
        List of normalized rewards.
    """
    rew_arr = np.asarray(rew_list, dtype=np.float32)
    norm = rew_arr / np.sqrt(self.rew_rms.var + self.eps)
    norm = np.clip(norm, -self.clip_rew, self.clip_rew)
    return norm.tolist()

def reset(self, *, seed: int | None = None, options: Dict[str, Any] | None = None):
    """
    Reset all environments and normalize observations.

    Args:
        seed: Optional base seed forwarded to the underlying vector env.
        options: Reset options forwarded to the underlying vector env.

    Returns:
        A pair ``(obs_list, info_list)``. Observations may be normalized.
    """
    obs_list, info_list = self.env.reset(seed=seed, options=options)

    if self.norm_obs and self.training:
        self.obs_rms.update(np.asarray(obs_list, dtype=np.float32))

    self.returns[:] = 0.0

    if self.norm_obs:
        obs_list = self._normalize_obs(obs_list)

    return obs_list, info_list

def reset_done(
    self, 
    dones: Union[List[bool], np.ndarray],
    *, 
    seed: int | None = None, 
    options: Dict[str, Any] | None = None
):
    """
    Reset only environments indicated by ``dones`` and normalize those observations.

    Args:
        dones: Boolean mask/list of length :attr:`n_envs`.
        seed: Optional base seed forwarded to the underlying vector env.
        options: Reset options forwarded to the underlying vector env.

    Returns:
        A tuple ``(reset_obs, reset_infos)`` as described by
        :meth:`VecEnvWrapperBase.reset_done`, with reset observations optionally
        normalized.
    """
    reset_obs, reset_infos = self.env.reset_done(
        dones, seed=seed, options=options
    )

    obs_arr = np.asarray(
        [o for o in reset_obs if o is not None],
        dtype=np.float32,
    ) if any(o is not None for o in reset_obs) else None

    if self.norm_obs and self.training and obs_arr is not None:
        self.obs_rms.update(obs_arr)

    for i, done in enumerate(dones):
        if done:
            self.returns[i] = 0.0

    if self.norm_obs:
        norm_reset_obs: List[Any] = list(reset_obs)
        # Only normalize indices that were reset
        for i, done in enumerate(dones):
            if done and reset_obs[i] is not None:
                o = np.asarray(reset_obs[i], dtype=np.float32)
                norm = (o - self.obs_rms.mean) / np.sqrt(self.obs_rms.var + self.eps)
                norm = np.clip(norm, -self.clip_obs, self.clip_obs)
                norm_reset_obs[i] = norm
        reset_obs = norm_reset_obs

    return reset_obs, reset_infos

def step(self, actions):
    """
    Step all environments and apply observation/reward normalization.

    Args:
        actions: Iterable of actions of length :attr:`n_envs`.

    Returns:
        A 5-tuple ``(obs_list, rew_list, term_list, trunc_list, infos)``, where
        observations and/or rewards may be normalized.
    """
    obs_list, rew_list, term_list, trunc_list, infos = self.env.step(actions)

    obs_arr = np.asarray(obs_list, dtype=np.float32)
    rew_arr = np.asarray(rew_list, dtype=np.float32)

    if self.norm_obs and self.training:
        self.obs_rms.update(obs_arr)

    if self.norm_rew:
        self.returns = self.returns * self.gamma + rew_arr
        if self.training:
            self.rew_rms.update(self.returns)

        rew_list = self._normalize_rew(rew_list)

    if self.norm_obs:
        obs_list = self._normalize_obs(obs_list)

    return obs_list, rew_list, term_list, trunc_list, infos