Skip to content

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

src.joint_graph.common.Graph dataclass

Persistence

src.joint_graph.serialise.save_graph(graph, path, *, protocol=pickle.HIGHEST_PROTOCOL)

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

src.joint_graph.serialise.load_graph(path)

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

Graph Augmentation and Induction

src.joint_graph.augment.apply_world_model_env_and_rewards(graph, *, model, encode_state_action_fn, predict_probs_bits_rewards_fn, action_size, n_agents, cells, n_channels, focal_idx=0, inplace=False, wm_state_bits_fn=None, eps=1e-09)

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(graph, *, opponent_policies, focal_idx=0, inplace=False, policy_state_fn=None)

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(graph, *, opponent_policy, focal_idx=0, inplace=False, policy_state_fn=None, base_action_size=None)

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(graph, *, focal_idx=0, action_size_self=None)

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(p)

src.joint_graph.policy_utils.make_policy_prob_fn(policy)

src.joint_graph.visualise.graph_to_dot(graph, *, precision=3, show_state_id=False)

src.joint_graph.visualise.export_dot(graph, path, *, precision=3, show_state_id=False)