Linear Temporal Logic (LTL) Safety Constraint¶

def reset(self):
    """Reset the safety monitor and underlying DFA-cost state."""
    self.safe = True
    self.step_cost = 0.0
    self.total_unsafe = 0.0
    self.cost_fn.reset()

def update(self, labels: Iterable[str]):
    """Update the DFA-cost state and safety flags.

    Args:
        labels: Iterable of atomic propositions true at the current step.
    """
    self.step_cost = self.cost_fn(labels)
    self.total_unsafe += float(self.step_cost >= 0.5)
    self.safe = self.safe and (not self.total_unsafe)

def get_automaton_state(self):
    """Return the current DFA state from the underlying DFA-cost object."""
    return self.cost_fn.automaton_state

def get_dfa(self):
    """Return the DFA used by the underlying DFA-cost object."""
    return self.cost_fn.dfa

def satisfied(self) -> bool:
    """Whether the episode remains safe so far."""
    return self.safe

def episode_metric(self) -> Dict[str, float]:
    """End-of-episode metrics.

    Returns:
        Dict containing:

        - ``"cum_unsafe"``: count of unsafe steps,
        - ``"satisfied"``: 1.0 if safe else 0.0.
    """
    return {
        "cum_unsafe": float(self.total_unsafe),
        "satisfied": float(self.satisfied()),
    }

def step_metric(self) -> Dict[str, float]:
    """Per-step metrics.

    Returns:
        Dict containing:

        - ``"cost"``: current step cost,
        - ``"violation"``: 1.0 if ``cost >= 0.5`` else 0.0.
    """
    return {
        "cost": self.step_cost,
        "violation": float(self.step_cost >= 0.5),
        "cost_done": float(self.step_cost >= 0.5),
    }

def _make_augmented_obs_space(
    self, orig: spaces.Space, obs_type: str
) -> spaces.Space:
    """Construct the augmented observation space.

    Args:
        orig: Original observation space of the wrapped environment.

    Returns:
        A new observation space that includes the automaton state.

    Raises:
        TypeError: If the observation space type is unsupported, or if a Box
            space is not 1-D.
    """
    if isinstance(orig, spaces.Discrete):
        n = int(orig.n)
        if obs_type == "discrete":
            return spaces.Discrete(n * self._num_automaton_states)
        if obs_type == "box":
            dim = n + self._num_automaton_states
            return spaces.Box(
                low=0.0,
                high=1.0,
                shape=(dim,),
                dtype=self._box_dtype,
            )
        if obs_type == "dict":
            return spaces.Dict(
                {
                    "orig": spaces.Discrete(n),
                    "automaton": spaces.Discrete(self._num_automaton_states),
                }
            )
        raise RuntimeError(f"Unhandled obs_type: {obs_type!r}")

    if isinstance(orig, spaces.Box):
        if obs_type == "discrete":
            raise TypeError(
                "Incompatible configuration: obs_type='discrete' requires a Discrete "
                "base observation space, but got Box."
            )
        if orig.shape is None or len(orig.shape) != 1:
            raise TypeError(
                f"LTLSafetyEnv only supports 1-D Box for augmentation; got shape {orig.shape}"
            )
        d = int(orig.shape[0])
        if obs_type == "box":
            low = np.concatenate(
                [
                    orig.low.astype(self._box_dtype, copy=False),
                    np.zeros(self._num_automaton_states, dtype=self._box_dtype),
                ]
            )
            high = np.concatenate(
                [
                    orig.high.astype(self._box_dtype, copy=False),
                    np.ones(self._num_automaton_states, dtype=self._box_dtype),
                ]
            )
            return spaces.Box(low=low, high=high, dtype=self._box_dtype)
        if obs_type == "dict":
            return spaces.Dict(
                {
                    "orig": orig,
                    "automaton": spaces.Discrete(self._num_automaton_states),
                }
            )
        raise RuntimeError(f"Unhandled obs_type: {obs_type!r}")

    if isinstance(orig, spaces.Dict):
        if obs_type == "discrete":
            raise TypeError(
                "Incompatible configuration: obs_type='discrete' requires a Discrete "
                "base observation space, but got Dict."
            )
        if obs_type == "box":
            raise TypeError(
                "Incompatible configuration: obs_type='box' is not supported when the "
                "base observation space is Dict (cannot flatten generically). "
                "Use obs_type='dict' instead."
            )
        new_spaces = dict(orig.spaces)
        new_spaces["automaton"] = spaces.Discrete(self._num_automaton_states)
        return spaces.Dict(new_spaces)

    raise TypeError(
        f"LTLSafetyEnv does not support base observation space type {type(orig).__name__}. "
        "Supported base spaces: Discrete, 1-D Box, Dict."
    )

def _one_hot(self, idx: int, dim: int) -> np.ndarray:
    """One-hot encode an index into a vector of length ``dim``.

    Args:
        idx: index to encomde.
        dim: length of one-hot encoding.

    Returns:
        A 1-D numpy array of shape ``(dim,)`` containing a
        one-hot encoding. If ``idx`` is out of range, returns the all-zeros vector.
    """
    enc = np.zeros(dim, dtype=self._box_dtype)
    if 0 <= int(idx) < dim:
        enc[int(idx)] = 1.0
    return enc

def _augment_obs(self, obs: Any) -> Any:
    """Augment a base observation with the current automaton state.

    Args:
        obs: Base observation returned by the wrapped environment.

    Returns:
        Augmented observation matching :attr:`observation_space`.

    Raises:
        TypeError: If the base observation does not match the expected type/shape
            implied by the observation space.
        RuntimeError: If the wrapper is in an unexpected observation-space state.
    """
    q_state = self._constraint.get_automaton_state()
    q_idx = int(self._automaton_states_idx[q_state])

    orig = self._orig_obs_space
    obs_type = self._obs_type

    if isinstance(orig, spaces.Discrete):
        if not isinstance(obs, (int, np.integer)):
            raise TypeError(
                f"Expected Discrete obs as int, got {type(obs).__name__}"
            )
        s = int(obs)
        if not (0 <= s < int(orig.n)):
            raise TypeError(
                f"Discrete obs out of range: got {s}, expected [0, {orig.n})"
            )
        if obs_type == "discrete":
            return int(orig.n) * q_idx + s
        if obs_type == "box":
            return np.concatenate(
                [
                    self._one_hot(s, int(orig.n)),
                    self._one_hot(q_idx, self._num_automaton_states),
                ],
                axis=0,
            )
        if obs_type == "dict":
            return {
                "orig": s,
                "automaton": q_idx,
            }
        raise RuntimeError(f"Unhandled obs_type: {obs_type!r}")
    if isinstance(orig, spaces.Box):
        if obs_type == "discrete":
            raise RuntimeError(
                "obs_type='discrete' with Box base should have been rejected."
            )
        arr = (
            obs
            if isinstance(obs, np.ndarray)
            else np.asarray(obs, dtype=self._box_dtype)
        )
        if arr.shape is None or arr.ndim != 1:
            raise TypeError(f"Expected 1-D Box observation, got shape {arr.shape}")
        arr = arr.astype(self._box_dtype, copy=False)
        if obs_type == "box":
            return np.concatenate(
                [arr, self._one_hot(q_idx, self._num_automaton_states)], axis=0
            )
        if obs_type == "dict":
            return {"orig": arr, "automaton": q_idx}
        raise RuntimeError(f"Unhandled obs_type: {obs_type!r}")
    if isinstance(orig, spaces.Dict):
        if obs_type != "dict":
            raise RuntimeError(
                "Only obs_type='dict' is supported for Dict base spaces."
            )
        if not isinstance(obs, dict):
            raise TypeError(f"Expected Dict obs as dict, got {type(obs).__name__}")
        out = dict(obs)
        out["automaton"] = q_idx
        return out

    raise RuntimeError(
        f"Unexpected base observation space type {type(orig).__name__}"
    )

def reset(self, *, seed: int | None = None, options: Dict[str, Any] | None = None):
    obs, info = self.env.reset(seed=seed, options=options)
    self._constraint.reset()
    labels = info.get("labels", set())
    self._constraint.update(labels)
    info["automaton_state"] = self._constraint.get_automaton_state()
    return self._augment_obs(obs), info

def step(self, action):
    obs, reward, terminated, truncated, info = self.env.step(action)
    labels = info.get("labels", set())
    self._constraint.update(labels)
    info["automaton_state"] = self._constraint.get_automaton_state()
    return self._augment_obs(obs), reward, terminated, truncated, info