Misc Wrappers¶

def _setup_cost_fn(self):
    """
    Create :attr:`shaped_cost_fn` if DFA-based constraints are available.

    If the underlying constraint exposes a :class:`~masa.common.ltl.DFACostFn`,
    constructs a :class:`~masa.common.ltl.ShapedCostFn`. Otherwise, uses a
    trivial zero-cost function.

    Returns:
        ``None``. This method sets :attr:`shaped_cost_fn` as a side effect.
    """
    if hasattr(self._constraint, "cost_fn") and isinstance(self._constraint.cost_fn, DFACostFn):
        self.shaped_cost_fn = ShapedCostFn(self._constraint.cost_fn.dfa, self.potential_fn, gamma=self._gamma)
    else:
        self.shaped_cost_fn = lambda q: 0.0

def _setup_potential_fn(self):
    """
    Construct the potential function :attr:`potential_fn` for shaping.

    For ``impl="none"``, :attr:`potential_fn` is identically zero.

    For ``impl="vi"``, a small fixed number of value-iteration steps are run
    over DFA states, treating accepting states as having a constant terminal
    value and propagating backward through reachable transitions.

    For ``impl="cycle"``, a reverse-graph BFS is used to find a "furthest"
    state from the accepting set and then compute distances to that target,
    yielding a shaping potential based on distance.

    Returns:
        ``None``. This method sets :attr:`potential_fn` and may store
        intermediate tables such as :attr:`V` or :attr:`dist_to_furthest`.

    Raises:
        AssertionError: If ``impl != "none"`` but no DFA-based cost function
        is available on the underlying constraint.
    """

    if self._impl != "none":
        assert hasattr(self._constraint, "cost_fn"), \
        ("RewardShapingWrapper requires env to implement a BaseConstraintEnv that exposes a cost_fn")
        assert isinstance(getattr(self._constraint, "cost_fn", None), DFACostFn), \
        ("RewardShapingWrapper requires env to implement a LTLSafetyEnv with cost_fn class: DFACostFn")

        dfa: DFA = self.env._constraint.cost_fn.dfa
    else:
        self.potential_fn = lambda q: 0.0

    if self._impl == "vi":

        VI_STEPS = 100
        GAMMA = 0.9
        self.V = {q: 0.0 for q in dfa.states}
        assert GAMMA <= self._gamma

        print("Reward shaping DFA ...")
        for i in tqdm(range(VI_STEPS)):
            diff = 0.0
            for u in dfa.states:
                V_u = self.V[u]
                self.V[u] = 1.0/(1.0 - GAMMA) if u in dfa.accepting else \
                    np.max([GAMMA * self.V[v] for v in dfa.edges[u].keys()])
                diff = max(diff, np.abs(V_u - self.V[u]))
            if diff < 1e-6:
                break

        self.potential_fn = lambda q: self.V[q]

    if self._impl == "cycle":

        print("Reward shaping DFA ...")
        edges_rev = {v: set() for v in dfa.states}
        for u in dfa.states:
            if u in dfa.edges:
                reachable_states = set(dfa.edges[u].keys())
                for v in reachable_states:
                    edges_rev[v].add(u)

        dist_to_accepting = {u: np.inf for u in dfa.states}
        queue = deque()

        for a in dfa.accepting:
            dist_to_accepting[a] = 0.0
            queue.append(a)

        max_dist = -1
        furthest_state = None

        while queue:
            current = queue.popleft()
            current_dist = dist_to_accepting[current]

            if current_dist > max_dist:
                max_dist = current_dist
                furthest_state = current

            for w in edges_rev.get(current, []):
                if dist_to_accepting[w] == np.inf:
                    dist_to_accepting[w] = current_dist + 1
                    queue.append(w)

        if furthest_state is None:
            furthest_state = dfa.initial

        self.dist_to_furthest = {u: np.inf for u in dfa.states}
        max_finite_dist = 0.0

        if furthest_state is not None:
            u_target = furthest_state
            self.dist_to_furthest[u_target] = 0.0
            queue = deque([u_target])

            while queue:
                current = queue.popleft()
                current_dist = self.dist_to_furthest[current]

                if current_dist > max_finite_dist:
                    max_finite_dist = current_dist

                for w in edges_rev.get(current, []):
                    if self.dist_to_furthest[w] == np.inf:
                        self.dist_to_furthest[w] = current_dist + 1
                        queue.append(w)

        replacement_value = max_finite_dist + 1.0

        for u in dfa.states:
            if self.dist_to_furthest[u] == np.inf:
                self.dist_to_furthest[u] = replacement_value

        self.potential_fn = lambda q: self.dist_to_furthest[q]

def reset(self, *, seed: int | None = None, options: Dict[str, Any] | None = None):
    """
    Reset the environment and initialize shaping state.

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

    Returns:
        A tuple ``(obs, info)`` from the underlying environment.

    Notes:
        This wrapper reads ``info["automaton_state"]`` to initialize the
        previous potential :attr:`_last_potential`. If the key is missing,
        it assumes DFA state ``0``.
    """
    obs, info = self.env.reset(seed=seed, options=options)
    automaton_state = info.get("automaton_state", 0)
    self._last_potential = self.potential_fn(automaton_state)
    return obs, info

def step(self, action: Any):
    """
    Step the environment and apply potential-based shaping to the step cost.

    Args:
        action: Action forwarded to the underlying environment.

    Returns:
        A 5-tuple ``(observation, reward, terminated, truncated, info)``.

    Side effects:
        Updates ``info["constraint"]["step"]["cost"]`` in-place with the shaped
        cost and updates :attr:`_last_potential`.

    Notes:
        If the underlying ``info`` does not contain constraint metrics,
        this method assumes an unshaped step cost of ``0.0`` and will
        still attempt to write back into ``info["constraint"]["step"]``.
    """
    observation, reward, terminated, truncated, info = self.env.step(action)
    cost = info["constraint"]["step"].get("cost", 0.0)
    automaton_state = info.get("automaton_state", 0)
    potential = self.potential_fn(automaton_state)
    info["constraint"]["step"]["cost"] = cost + self._gamma * potential - self._last_potential
    self._last_potential = potential
    return observation, reward, terminated, truncated, info

def _normalize_obs(self, obs: np.ndarray) -> np.ndarray:
    """
    Normalize (and clip) a single observation.

    Args:
        obs: Raw observation.

    Returns:
        Normalized observation.
    """
    return np.clip(
        (obs - self.obs_rms.mean) / np.sqrt(self.obs_rms.var + self.eps),
        -self.clip_obs,
        self.clip_obs
    )

def _normalize_rew(self, rew: float) -> float:
    """
    Normalize (and clip) a single reward.

    Args:
        rew: Raw reward.

    Returns:
        Normalized reward.
    """
    return np.clip(
        rew / np.sqrt(self.rew_rms.var + self.eps),
        -self.clip_rew,
        self.clip_rew,
    )

def reset(self, *, seed: int | None = None, options: Dict[str, Any] | None = None):
    """
    Reset the environment and (optionally) update normalization statistics.

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

    Returns:
        A tuple ``(obs, info)`` where ``obs`` may be normalized.
    """
    obs, info = self.env.reset(seed=seed, options=options)

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

    if self.norm_rew and self.training:
        self.returns[:] = 0.0

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

    return obs, info

def step(self, action):
    """
    Step the environment and apply observation/reward normalization.

    Args:
        action: Action forwarded to the underlying environment.

    Returns:
        A 5-tuple ``(obs, rew, terminated, truncated, info)`` where ``obs`` and/or
        ``rew`` may be normalized.
    """
    obs, rew, term, trunc, info = self.env.step(action)

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

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

        rew = self._normalize_rew(rew)

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

    return obs, rew, term, trunc, info

@staticmethod
def _one_hot_scalar(idx: int, n: int) -> np.ndarray:
    """
    One-hot encode an integer index.

    Args:
        idx: Index in ``{0, 1, ..., n-1}``.
        n: Vector length.

    Returns:
        A float32 vector ``v`` with ``v[idx] = 1`` and zeros elsewhere.
    """
    one_hot = np.zeros(n, dtype=np.float32)
    one_hot[idx] = 1.0
    return one_hot

def _get_obs(self, obs: Union[int, Dict[str, Any], np.ndarray]) -> np.ndarray:
    """
    Transform an observation according to the wrapper's configured mode.

    Args:
        obs: Raw observation.

    Returns:
        One-hot encoded observation (or dict containing one-hot fields) when applicable,
        otherwise the original observation.
    """
    if self._mode == "discrete":
        # Original obs_space is Discrete; obs is an int-like
        idx = int(obs)
        n = self._orig_obs_space.n
        return self._one_hot_scalar(idx, n)

    elif self._mode == "dict":
        assert isinstance(obs, dict), (
            f"Expected dict observation for Dict space, got {type(obs)}"
        )

        new_obs: Dict[str, Any] = {}
        for key, subspace in self._orig_obs_space.spaces.items():
            value = obs[key]

            if isinstance(subspace, spaces.Discrete):
                idx = int(value)
                new_obs[key] = self._one_hot_scalar(idx, subspace.n)
            else:
                # Leave non-Discrete parts unchanged
                new_obs[key] = value

        return new_obs
    else: 
        # pass
        return obs

def _get_obs(self, obs: Dict[str, Any]) -> np.ndarray:
    """
    Flatten a Dict observation into a 1D vector.

    Args:
        obs: Dict observation keyed the same way as the original Dict space.

    Returns:
        A 1D float32 array created by concatenating flattened sub-observations.

    Raises:
        TypeError: If any subspace is not a Box.

    Notes:
        Although the constructor supports Discrete subspaces when building bounds,
        this implementation currently enforces Box-only subspaces at runtime.
    """
    assert isinstance(obs, dict), (
            f"Expected dict observation for Dict space, got {type(obs)}"
        )

    parts = []
    for key in sorted(self._orig_obs_space.spaces.keys()):
        subspace = self._orig_obs_space.spaces[key]
        value = obs[key]

        if not isinstance(subspace, spaces.Box):
            raise TypeError(
                f"FlattenDictObsWrapper only supports Box subspaces, "
                f"got {type(subspace)} for key '{key}'"
            )

        arr = np.asarray(value, dtype=np.float32).reshape(-1)
        parts.append(arr)

    return np.concatenate(parts).astype(np.float32)

def _capture_frame(self):
    assert self.recording, "Cannot capture a frame, recording wasn't started."

    frame = self.env.render()
    if isinstance(frame, list):
        if len(frame) == 0:
            return
        self.render_history += frame
        frame = frame[-1]

    if isinstance(frame, np.ndarray):
        self.recorded_frames.append(frame)
    else:
        self.stop_recording()
        gymnasium.logger.warn(
            f"Recording stopped: expected type of frame returned by render to be a numpy array, got {type(frame)}."
        )

def reset(
    self, seed: int | None = None, options: dict | None = None
) -> tuple[dict[AgentID, ObsType], dict[AgentID, dict]]:
    obs, info = self.env.reset(seed=seed, options=options)
    self.episode_id += 1

    if self.recording and self.video_length == float("inf"):
        self.stop_recording()

    if self.episode_trigger and self.episode_trigger(self.episode_id):
        self.start_recording(f"{self.name_prefix}-episode-{self.episode_id}")
    if self.recording:
        self._capture_frame()
        if len(self.recorded_frames) > self.video_length:
            self.stop_recording()

    return obs, info

def step(
    self, actions: dict[AgentID, ActionType]
) -> tuple[
    dict[AgentID, ObsType],
    dict[AgentID, float],
    dict[AgentID, bool],
    dict[AgentID, bool],
    dict[AgentID, dict],
]:
    obs, rew, terminated, truncated, info = self.env.step(actions)
    self.step_id += 1

    if self.step_trigger and self.step_trigger(self.step_id):
        self.start_recording(f"{self.name_prefix}-step-{self.step_id}")
    if self.recording:
        self._capture_frame()
        if len(self.recorded_frames) > self.video_length:
            self.stop_recording()

    return obs, rew, terminated, truncated, info

def render(self):
    render_out = self.env.render()
    if self.recording and isinstance(render_out, list):
        self.recorded_frames += render_out

    if len(self.render_history) > 0:
        tmp_history = self.render_history
        self.render_history = []
        frames = render_out if isinstance(render_out, list) else [render_out]
        return tmp_history + frames
    return render_out

def close(self):
    super().close()
    if self.recording:
        self.stop_recording()

def start_recording(self, video_name: str):
    if self.recording:
        self.stop_recording()

    self.recording = True
    self._video_name = video_name

def stop_recording(self):
    assert self.recording, "stop_recording was called, but no recording was started"

    if len(self.recorded_frames) == 0:
        gymnasium.logger.warn("Ignored saving a video as there were zero frames to save.")
    else:
        try:
            from moviepy.video.io.ImageSequenceClip import ImageSequenceClip
        except ImportError as e:
            raise DependencyNotInstalled(
                'MoviePy is not installed, run `pip install "moviepy>=2.2.1,<3.0.0"`'
            ) from e

        clip = ImageSequenceClip(self.recorded_frames, fps=self.frames_per_sec)
        moviepy_logger = None if self.disable_logger else "bar"
        path = os.path.join(self.video_folder, f"{self._video_name}.mp4")
        clip.write_videofile(path, logger=moviepy_logger)

    self.recorded_frames = []
    self.recording = False
    self._video_name = None

    if self.gc_trigger and self.gc_trigger(self.episode_id):
        gc.collect()

def __del__(self):
    if len(getattr(self, "recorded_frames", [])) > 0:
        gymnasium.logger.warn("Unable to save last video! Did you call close()?")