Joint Graphs API¶

Joint graph APIs define the shared graph representation used by world modelling, opponent modelling, and shielding.

Data Structures¶

src.joint_graph.common.Edge `dataclass` ¶

Edge(
    src: StateKey,
    dst: StateKey,
    joint_action: JointAction,
    env_prob: float,
    opp_prob: float,
    prob: float,
    rewards: ndarray,
)

src.joint_graph.common.Graph `dataclass` ¶

Graph(
    nodes: Dict[StateKey, ndarray],
    edges: List[Edge],
    start_keys: List[StateKey],
)

Persistence¶

src.joint_graph.serialise.save_graph ¶

save_graph(
    graph: Graph,
    path: str | Path | BinaryIO,
    *,
    protocol: int = pickle.HIGHEST_PROTOCOL,
) -> None

Persist a built graph to disk or a binary file handle.

src.joint_graph.serialise.load_graph ¶

load_graph(path: str | Path | BinaryIO) -> Graph

Load a previously saved graph from disk or a binary file handle.

Exact Legal Graphs¶

src.joint_graph.exact.build_legal_joint_graph ¶

build_legal_joint_graph(env, **kwargs)

src.joint_graph.exact.resolve_legal_joint_graph_builder ¶

resolve_legal_joint_graph_builder(env)

Graph Augmentation and Induction¶

src.joint_graph.augment.apply_world_model_env_and_rewards ¶

apply_world_model_env_and_rewards(
    graph: Graph,
    *,
    model,
    encode_state_action_fn: Callable[..., ndarray],
    predict_probs_bits_rewards_fn: Callable[
        ..., Tuple[ndarray, ndarray, ndarray]
    ],
    action_size: int,
    n_agents: int,
    cells: int,
    n_channels: int,
    focal_idx: int = 0,
    inplace: bool = False,
    wm_state_bits_fn: Optional[
        Callable[[ndarray], ndarray]
    ] = None,
    eps: float = 1e-09,
) -> Graph

Uses your world model to fill

env_prob on each edge (normalized over existing dst candidates in the graph for each (src, joint_action))
rewards on each edge (same rewards for all dst in a (src, joint_action) group, because your MLP reward head is conditioned on (s,a), not (s,a,s'))

World-model env_prob is computed as an (independent-bit) likelihood score: log P(s') = sum_i [b_i log p_i + (1-b_i) log(1-p_i)] then normalized across candidate dsts.

src.joint_graph.augment.apply_opponent_policies ¶

apply_opponent_policies(
    graph: Graph,
    *,
    opponent_policies: Sequence,
    focal_idx: int = 0,
    inplace: bool = False,
    policy_state_fn: Optional[
        Callable[[ndarray], ndarray]
    ] = None,
) -> Graph

Fills Edge.opp_prob for each (src, joint_action) based on opponent policies, and updates Edge.prob if env_prob is known.

opponent_policies must be in agent-index order excluding focal_idx

[policy for agent 0..N-1 if i != focal_idx]

src.joint_graph.augment.apply_iop_policies ¶

apply_iop_policies(
    graph: Graph,
    *,
    opponent_policy,
    focal_idx: int = 0,
    inplace: bool = False,
    policy_state_fn: Optional[
        Callable[[ndarray], ndarray]
    ] = None,
    base_action_size: Optional[int] = None,
) -> Graph

Option A: treat all non-focal agents as ONE JOINT opponent.

Fills Edge.opp_prob for each (src, joint_action) by: opp_prob = π_joint(a_-i_tuple | s)

Where a_-i_tuple is taken from joint_action excluding focal_idx, and encoded as a base-A index in opponent-agent order: opponent_indices = [0..N-1 excluding focal_idx] (in increasing env index order)

opponent_policy must output a vector of length A**(N-1).

src.joint_graph.induce.induce_focal_mdp ¶

induce_focal_mdp(
    graph: Any,
    *,
    focal_idx: int = 0,
    action_size_self: Optional[int] = None,
) -> IndGraph

Convert a joint-action graph into a focal-action induced MDP.

Requirements on graph: - graph.nodes: dict[state_key -> bits/array-like] - graph.edges: iterable of edges with fields: .src, .dst, .joint_action (tuple[int,...]) and either: .prob or: .env_prob (+ optional .opp_prob)

Output

src.shield.induced.types.Graph with edges keyed by a_self only (marginalized over opponents)

Policy and Visualisation Helpers¶

src.joint_graph.policy_utils.normalize_probs ¶

normalize_probs(p: ndarray) -> np.ndarray

src.joint_graph.policy_utils.make_policy_prob_fn ¶

make_policy_prob_fn(policy)

src.joint_graph.visualise.graph_to_dot ¶

graph_to_dot(
    graph: Graph,
    *,
    precision: int = 3,
    show_state_id: bool = False,
) -> str

src.joint_graph.visualise.export_dot ¶

export_dot(
    graph: Graph,
    path: str,
    *,
    precision: int = 3,
    show_state_id: bool = False,
) -> str