Model Checking¶

def _build_vec_label_fn(self) -> np.ndarray:
    """Build :attr:`vec_label_fn` from :attr:`label_fn`.

    Returns:
        Float64 array ``vec`` with shape ``(n_atoms, n_states)`` where
        ``vec[i, s] == 1.0`` iff atom ``i`` holds in state ``s``.
    """
    n_atoms = len(self.atomic_predicates)
    vec = np.zeros((n_atoms, self.n_states), dtype=np.float64)
    for s in range(self.n_states):
        labels = self.label_fn(s) 
        for atom, idx in self.atom_dict.items():
            vec[idx, s] = 1.0 if atom in labels else 0.0
    return vec

def update_kernel(
    self,
    transition_matrix: Optional[np.ndarray] = None,
    successor_states: Optional[np.ndarray] = None,
    probabilities: Optional[np.ndarray] = None,
):
    """Update the stored transition representation in-place.

    Exactly one update mode should be used:

    - Dense update: provide ``transition_matrix`` with the same shape as the
      original.
    - Compact update: provide both ``successor_states`` and ``probabilities``
      with shapes consistent with the original compact representation.

    Args:
        transition_matrix: New dense MDP kernel ``(S, S, A)``.
        successor_states: New successor ids ``(K, S)``.
        probabilities: New successor probabilities ``(K, S, A)``.

    Raises:
        AssertionError: If shapes do not match the originally configured
            representation.
    """
    if transition_matrix is not None:
        assert self.transition_matrix is not None
        tm = np.asarray(transition_matrix, dtype=np.float64)
        assert tm.ndim == 3
        S, S2, A = tm.shape
        assert S == S2
        assert tm.shape == self.transition_matrix.shape
        self.transition_matrix = tm
        self.successor_states = None
        self.probabilities = None
    else:
        assert self.successor_states is not None and self.probabilities is not None
        assert successor_states is not None and probabilities is not None
        succ = np.asarray(successor_states, dtype=np.int64)
        probs = np.asarray(probabilities, dtype=np.float64)
        assert succ.ndim == 2
        assert probs.ndim == 3
        K, S = succ.shape
        K2, S2, A = probs.shape
        assert K == K2 and S == S2
        self.transition_matrix = None
        assert S == self.n_states and A == self.n_actions
        self.successor_states = succ
        self.probabilities = probs

def check_state(
    self,
    key: jax.Array, 
    policy: np.array
) -> np.ndarray:
    """Evaluate the stored formula on the policy-induced Markov chain.

    Args:
        key: Unused PRNG key (kept for API symmetry with
            :class:`StatisticalModelChecker`).
        policy: Stochastic policy probabilities, either:
            - shape ``(n_actions, n_states)`` (action-major), or
            - shape ``(n_states, n_actions)`` (state-major).

    Returns:
        Float64 array of shape ``(n_states,)`` in ``{0.0, 1.0}``.

    Raises:
        ValueError: If ``policy`` has an unexpected shape.
    """

    policy = np.asarray(policy, dtype=np.float64)

    tm = self.transition_matrix

    if policy.shape == (A, S):
        policy_sa = policy.T
        m_pi = np.einsum('ija,ai->ij', tm, policy)
    elif policy.shape == (S, A):
        policy_sa = policy
        m_pi = np.einsum('ija,ia->ij', tm, policy)
    else:
        raise ValueError(
            f"Unexpected policy shape {policy.shape}; expected "
            f"(n_actions, n_states) or (n_states, n_actions)"
        )

    if self.mode == "full":
        tm = self.transition_matrix
        m_pi = np.einsum("nsa,sa->ns", tm, policy_sa).astype(np.float64)
        kernel: Kernel = m_pi

    else:
        succ = self.successor_states
        probs = self.probabilities
        p_pi = np.einsum("ksa,sa->ks", probs, policy_sa).astype(np.float64)
        kernel = (succ, p_pi)


    return self.formula.sat(kernel, self.vec_label_fn, self.atom_dict)

def check_state_action(
    self,
    key: jax.Array,
    policy: np.ndarray,
) -> np.ndarray:
    """Compute a state-action satisfaction value for the stored formula.

    This method:
    1) Builds the policy-induced Markov chain kernel,
    2) Computes the formula sequence up to bound ``B = formula.bound``,
    3) Uses the vector at time ``max(B-1, 0)`` as a value function,
    4) Computes one-step expectations under each action to produce ``Q``.

    Args:
        key: Unused PRNG key (kept for API symmetry with
            :class:`StatisticalModelChecker`).
        policy: Stochastic policy probabilities, either ``(A, S)`` or ``(S, A)``.

    Returns:
        Float64 array ``Q`` of shape ``(n_states, n_actions)`` where
        ``Q[s, a]`` is the expected satisfaction value after taking action
        ``a`` in state ``s`` and then following ``policy``.

    Raises:
        ValueError: If ``policy`` has an unexpected shape.
    """

    policy = np.asarray(policy, dtype=np.float64)

    if policy.shape == (self.n_actions, self.n_states):
        policy_sa = policy.T
    elif policy.shape == (self.n_states, self.n_actions):
        policy_sa = policy
    else:
        raise ValueError(
            f"Unexpected policy shape {policy.shape}; expected "
            f"(n_actions, n_states) or (n_states, n_actions)"
        )

    if self.mode == "full":
        tm = self.transition_matrix
        m_pi = np.einsum('nsa,sa->ns', tm, policy_sa)
        kernel: Kernel = m_pi
    else:
        succ = self.successor_states
        probs = self.probabilities
        p_pi = np.einsum("ksa,sa->ks", probs, policy_sa).astype(np.float64)
        kernel = (succ, p_pi)

    B = self.formula.bound
    seq = self.formula._prob_seq(kernel, self.vec_label_fn, self.atom_dict, max_k=B)
    V_Bm1 = seq[max(B - 1, 0)]

    if self.mode == "full":
        tm = self.transition_matrix
        Q = np.einsum("nsa,n->sa", tm, V_Bm1).astype(np.float64)
    else:
        succ = self.successor_states
        probs = self.probabilities
        V_succ = V_Bm1[succ]
        Q = np.sum(probs * V_succ[:, :, None], axis=0).astype(np.float64)

    return Q

def check_state(
    self,
    key: jax.Array, 
    policy: np.array,
    state: int,
    num_samples: int,
) -> np.ndarray:
    r"""Estimate whether ``state`` satisfies the formula under ``policy``.

    For probabilistic temporal formulas, this samples ``num_samples`` paths of
    length ``max_steps = max(1, formula.bound)`` and estimates

    .. math::

       \hat{p} = \frac{1}{N}\sum_{i=1}^N \mathbf{1}\{\pi_i \models \varphi\}.

    Args:
        key: PRNG key used for sampling.
        policy: Stochastic policy probabilities, either ``(A, S)`` or ``(S, A)``.
        state: Start state index.
        num_samples: Number of trajectories to sample.

    Returns:
        Scalar float64 JAX array equal to ``1.0`` if ``\hat{p} >= p`` else
        ``0.0``, where ``p`` is the formula’s probability threshold.
    """
    if not self._is_probabilistic_formula(self.formula):
        val = self._eval_state_formula_python(self.formula, state)
        return jnp.array(float(val), dtype=jnp.float64)

    start_state = int(state)
    num_samples = int(num_samples)
    max_steps = max(1, int(self.formula.bound))

    policy_sa = self._prepare_policy_probs_jax(
        self.n_states, self.n_actions, policy
    )

    prob_threshold = self._get_formula_prob(self.formula)

    if self.mode == "full":
        tm = self.transition_matrix
        m_pi = np.einsum("nsa,sa->ns", tm, policy_sa).astype(np.float64)
        m_pi_j = jnp.asarray(m_pi, dtype=jnp.float64)

        p_hat = self._estimate_prob_dense(
            key=key,
            start_state=start_state,
            num_samples=num_samples,
            max_steps=max_steps,
            m_first=m_pi_j,
            m_rest=m_pi_j,
            formula=self.formula,
            vec_labels=self.vec_label_fn_jax,
            atom_dict=self.atom_dict
        )
    else:
        assert self.successor_states is not None and self.probabilities is not None
        succ = self.successor_states
        probs = self.probabilities

        p_pi = np.einsum("ksa,sa->ks", probs, policy_sa).astype(np.float64)
        p_pi_j = jnp.asarray(p_pi, dtype=jnp.float64)

        p_hat = self._estimate_prob_compact(
            key=key,
            start_state=start_state,
            num_samples=num_samples,
            max_steps=max_steps,
            succ=succ_j,
            p_first=p_pi_j,
            p_rest=p_pi_j,
            formula=self.formula,
            vec_labels=self.vec_label_fn_jax,
            atom_dict=self.atom_dict
        )

    return jnp.where(p_hat >= prob_threshold, 1.0, 0.0).astype(jnp.float64)

def check_state_action(
    self,
    key: jax.Array,
    policy: np.ndarray,
    state: int,
    action: int,
    num_samples: int,
) -> np.ndarray:
    r"""Estimate satisfaction when forcing the first action, then following ``policy``.

    The first transition uses the forced action ``action`` and subsequent
    steps follow the policy-induced kernel.

    Args:
        key: PRNG key used for sampling.
        policy: Stochastic policy probabilities, either ``(A, S)`` or ``(S, A)``.
        state: Start state index.
        action: Forced first action index.
        num_samples: Number of trajectories to sample.

    Returns:
        Scalar float64 JAX array equal to ``1.0`` if the estimated probability
        meets the formula's threshold, else ``0.0``.
    """
    if not self._is_probabilistic_formula(self.formula):
        val = self._eval_state_formula_python(self.formula, state)
        return jnp.array(float(val), dtype=jnp.float64)

    start_state = int(state)
    forced_action = int(action)
    num_samples = int(num_samples)
    max_steps = max(1, int(self.formula.bound))

    policy_sa = self._prepare_policy_probs_jax(
        self.n_states, self.n_actions, policy
    )

    prob_threshold = self._get_formula_prob(self.formula)

    if self.mode == "full":
        tm = self.transition_matrix
        m_a = tm[:, :, action]
        m_pi = np.einsum("nsa,sa->ns", tm, policy_sa).astype(np.float64)

        m_a_j = jnp.asarray(m_a, dtype=jnp.float64)
        m_pi_j = jnp.asarray(m_pi, dtype=jnp.float64)

        p_hat = self._estimate_prob_dense(
            key=key,
            start_state=start_state,
            num_samples=num_samples,
            max_steps=max_steps,
            m_first=m_a_j,
            m_rest=m_pi_j,
            formula=self.formula,
            vec_labels=self.vec_label_fn_jax,
            atom_dict=self.atom_dict
        )
    else:
        assert self.successor_states is not None and self.probabilities is not None
        succ = self.successor_states
        probs = self.probabilities

        p_a = probs[:, :, action]
        p_pi = np.einsum("ksa,sa->ks", probs, policy_sa).astype(np.float64)

        succ_j = jnp.asarray(succ, dtype=jnp.int64)
        p_a_j = jnp.asarray(p_a, dtype=jnp.float64)
        p_pi_j = jnp.asarray(p_pi, dtype=jnp.float64)

        p_hat = self._estimate_prob_compact(
            key=key,
            start_state=start_state,
            num_samples=num_samples,
            max_steps=max_steps,
            succ=succ_j,
            p_first=p_a_j,
            p_rest=p_pi_j,
            formula=self.formula,
            vec_labels=self.vec_label_fn_jax,
            atom_dict=self.atom_dict
        )

    return jnp.where(p_hat >= prob_threshold, 1.0, 0.0).astype(jnp.float64)

def _is_probabilistic_formula(self, formula: BoundedPCTLFormula) -> bool:
    """Return True if ``formula`` is a probabilistic path operator."""
    return isinstance(formula, (Next, Until, Eventually, Always))

def _get_formula_prob(self, formula: BoundedPCTLFormula) -> float:
    """Extract probability threshold from a probabilistic formula.

    Args:
        formula: A probabilistic formula (``Next``, ``Until``, ``Eventually``, ``Always``).

    Returns:
        Probability threshold.

    Raises:
        ValueError: If ``formula`` is not a probabilistic operator.
    """
    if isinstance(formula, (Next, Until, Eventually, Always)):
        return float(formula.prob)
    raise ValueError(
        "Trying to get probability threshold from non-probabilistic formula."
    )

def _eval_state_formula_python(self, formula: BoundedPCTLFormula, state: int) -> bool:
    """Evaluate a pure state formula in Python.

    This supports only boolean (non-probabilistic) formulas. Nested
    probabilistic operators are explicitly rejected.

    Args:
        formula: Formula to evaluate.
        state: State index.

    Returns:
        Boolean satisfaction.

    Raises:
        NotImplementedError: If nested probabilistic operators are encountered.
        TypeError: If an unsupported formula type is encountered.
    """
    if isinstance(formula, Truth):
        return True
    if isinstance(formula, Atom):
        idx = self.atom_dict[formula.atom]
        return bool(self.vec_label_fn[idx, state] > 0.5)
    if isinstance(formula, Neg):
        return not self._eval_state_formula_python(formula.subformula, state)
    if isinstance(formula, And):
        return (
            self._eval_state_formula_python(formula.subformula_1, state)
            and self._eval_state_formula_python(formula.subformula_2, state)
        )
    if isinstance(formula, Or):
        return (
            self._eval_state_formula_python(formula.subformula_1, state)
            or self._eval_state_formula_python(formula.subformula_2, state)
        )
    if isinstance(formula, Implies):
        left = self._eval_state_formula_python(formula.subformula_1, state)
        right = self._eval_state_formula_python(formula.subformula_2, state)
        return (not left) or right
    if isinstance(formula, (Next, Until, Eventually, Always)):
        raise NotImplementedError(
            "Nested probabilistic operators in state formulas are not "
            "supported in StatisticalModelChecker."
        )
    raise TypeError(f"Unsupported formula type in Python state eval: {type(formula)}")

@staticmethod
@partial(jit, static_argnames=["n_states", "n_actions"])
def _prepare_policy_probs_jax(
    n_states: int, n_actions: int, policy: jnp.ndarray
) -> jnp.ndarray:
    """Normalize policy shape for JAX computation.

    Args:
        n_states: Number of states.
        n_actions: Number of actions.
        policy: Policy probabilities as either ``(n_actions, n_states)`` or
            ``(n_states, n_actions)``.

    Returns:
        Policy as a ``(n_states, n_actions)`` JAX array.

    Raises:
        ValueError: If ``policy`` has an incompatible shape.
    """
    policy = jnp.asarray(policy, dtype=jnp.float64)
    if policy.shape == (n_actions, n_states):
        policy = policy.T
    elif policy.shape == (n_states, n_actions):
        pass
    else:
        raise ValueError(
            f"Policy shape {policy.shape} incompatible with "
            f"(n_states, n_actions)=({n_states}, {n_actions})"
        )
    return policy

@staticmethod
def _eval_state_formula_jax(
    state_idx: jnp.ndarray, 
    formula: BoundedPCTLFormula,
    vec_labels: jnp.ndarray, 
    atom_dict: Dict[str, int], 
) -> jnp.ndarray:
    """Evaluate a pure state formula in JAX.

    Args:
        state_idx: Scalar state index.
        formula: State formula (must not contain probabilistic operators).
        vec_labels: Vectorized labels ``(n_atoms, n_states)`` as a JAX array.
        atom_dict: Mapping from atom name to row index in ``vec_labels``.

    Returns:
        JAX boolean scalar indicating satisfaction.

    Raises:
        NotImplementedError: If nested probabilistic operators are encountered.
        TypeError: If an unsupported formula type is encountered.
    """

    def rec(f: BoundedPCTLFormula) -> jnp.ndarray:
        if isinstance(f, Truth):
            return jnp.bool_(True)
        if isinstance(f, Atom):
            idx = atom_dict[f.atom]
            return vec_labels[idx, state_idx] > 0.5
        if isinstance(f, Neg):
            return jnp.logical_not(rec(f.subformula))
        if isinstance(f, And):
            return jnp.logical_and(
                rec(f.subformula_1),
                rec(f.subformula_2),
            )
        if isinstance(f, Or):
            return jnp.logical_or(
                rec(f.subformula_1),
                rec(f.subformula_2),
            )
        if isinstance(f, Implies):
            left = rec(f.subformula_1)
            right = rec(f.subformula_2)
            return jnp.logical_or(jnp.logical_not(left), right)
        if isinstance(f, (Next, Until, Eventually, Always)):
            raise NotImplementedError(
                "Nested probabilistic operators in state formulas "
                "are not supported in JAX SMC."
            )
        raise TypeError(f"Unsupported formula type in JAX state eval: {type(f)}")

    return rec(formula)

@staticmethod
def _path_satisfies_single(
    states_1d: jnp.ndarray,
    formula: BoundedPCTLFormula,
    vec_labels: jnp.ndarray, 
    atom_dict: Dict[str, int], 
) -> jnp.ndarray:
    """Evaluate whether a single sampled path satisfies the formula.

    Args:
        states_1d: A single trajectory of state indices with shape ``(T+1,)``.
        formula: The (bounded) formula to check.
        vec_labels: Vectorized labels ``(n_atoms, n_states)`` as a JAX array.
        atom_dict: Mapping from atom name to row index.

    Returns:
        JAX boolean scalar: True iff the trajectory satisfies ``formula``.
    """
    eval_state = StatisticalModelChecker._eval_state_formula_jax

    if isinstance(formula, Next):
        s1 = states_1d[1]
        return eval_state(
            s1, formula.subformula, vec_labels, atom_dict
        )

    if isinstance(formula, Until):
        B = formula.bound
        idxs = jnp.arange(B + 1)
        s_seq = states_1d[idxs]
        sat1 = jax.vmap(lambda s: eval_state(
            s, formula.subformula_1, vec_labels, atom_dict
        ))(s_seq)
        sat2 = jax.vmap(lambda s: eval_state(
            s, formula.subformula_2, vec_labels, atom_dict
        ))(s_seq)

        def scan_body(carry, x):
            new_carry = jnp.logical_and(carry, x)
            return new_carry, new_carry

        init = jnp.bool_(True)
        _, prefix_tail = jax.lax.scan(scan_body, init, sat1[:-1])
        prefix_all = jnp.concatenate(
            [jnp.array([True], dtype=bool), prefix_tail], axis=0
        )
        cond_k = jnp.logical_and(sat2, prefix_all)
        return jnp.any(cond_k)

    if isinstance(formula, Eventually):
        B = formula.bound_param
        idxs = jnp.arange(B + 1)
        s_seq = states_1d[idxs]
        sat = jax.vmap(lambda s: eval_state(
            s, formula.subformula, vec_labels, atom_dict
        ))(s_seq)
        return jnp.any(sat)

    if isinstance(formula, Always):
        B = formula.bound_param
        idxs = jnp.arange(B + 1)
        s_seq = states_1d[idxs]
        sat = jax.vmap(lambda s: eval_state(
            s, formula.subformula, vec_labels, atom_dict
        ))(s_seq)
        return jnp.all(sat)

    return eval_state(states_1d[0], formula, vec_labels, atom_dict)

@staticmethod
def _estimate_prob_dense(
    key: jax.Array,
    start_state: int,
    num_samples: int,
    max_steps: int,
    m_first: jnp.ndarray,
    m_rest: jnp.ndarray,
    formula: BoundedPCTLFormula,
    vec_labels: jnp.ndarray,
    atom_dict: Dict[str, int],
) -> jnp.ndarray:
    """Estimate satisfaction probability by sampling from a dense kernel.

    Args:
        key: PRNG key.
        start_state: Initial state index.
        num_samples: Number of trajectories to sample.
        max_steps: Trajectory horizon (number of transitions).
        m_first: Transition matrix for the first step ``(n_states, n_states)``.
        m_rest: Transition matrix for subsequent steps ``(n_states, n_states)``.
        formula: Formula to check along sampled paths.
        vec_labels: Vectorized labels ``(n_atoms, n_states)`` as a JAX array.
        atom_dict: Atom-to-index mapping.

    Returns:
        Scalar float64 JAX array: estimated probability ``p_hat``.
    """
    states_batch = StatisticalModelChecker._sample_trajectories_dense_jax(
        key=key,
        m_first=m_first,
        m_rest=m_rest,
        start_state=jnp.asarray(start_state, dtype=jnp.float64),
        max_steps=max_steps,
        num_samples=num_samples,
    )

    path_satisfies_batch = jax.vmap(
        lambda s: StatisticalModelChecker._path_satisfies_single(
            s, formula, vec_labels, atom_dict
        ),
        in_axes=0,
    )

    sat_batch = path_satisfies_batch(states_batch)  # (num_samples,)
    return jnp.mean(sat_batch.astype(jnp.float64))

@staticmethod
def _estimate_prob_compact(
    key: jax.Array,
    start_state: int,
    num_samples: int,
    max_steps: int,
    succ: jnp.ndarray,
    p_first: jnp.ndarray,
    p_rest: jnp.ndarray,
    formula: BoundedPCTLFormula,
    vec_labels: jnp.ndarray,
    atom_dict: Dict[str, int],
) -> jnp.ndarray:
    """Estimate satisfaction probability by sampling from a compact kernel.

    Args:
        key: PRNG key.
        start_state: Initial state index.
        num_samples: Number of trajectories to sample.
        max_steps: Trajectory horizon (number of transitions).
        succ: Successor matrix ``(K, n_states)``.
        p_first: Probabilities for the first step ``(K, n_states)``.
        p_rest: Probabilities for subsequent steps ``(K, n_states)``.
        formula: Formula to check along sampled paths.
        vec_labels: Vectorized labels ``(n_atoms, n_states)`` as a JAX array.
        atom_dict: Atom-to-index mapping.

    Returns:
        Scalar float64 JAX array: estimated probability ``p_hat``.
    """
    states_batch = StatisticalModelChecker._sample_trajectories_compact_jax(
        key=key,
        succ=succ,
        p_first=p_first,
        p_rest=p_rest,
        start_state=jnp.asarray(start_state, dtype=jnp.float64),
        max_steps=max_steps,
        num_samples=num_samples,
    )

    path_satisfies_batch = jax.vmap(
        lambda s: StatisticalModelChecker._path_satisfies_single(
            s, formula, vec_labels, atom_dict
        ),
        in_axes=0,
    )

    sat_batch = path_satisfies_batch(states_batch)  # (num_samples,)
    return jnp.mean(sat_batch.astype(jnp.float64))