File size: 7,323 Bytes
d3a24e0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 | """State memory graph — Module C.
A hash-deduplicated directed graph of observed game states. Nodes are
hashed frames; edges are actions taken. Tracks visit counts for novelty
detection and supports backtracking/loop avoidance.
Inspired by Blind Squirrel's graph-based exploration, but integrated
with the learned perception encoder for efficient similarity checks.
"""
from __future__ import annotations
import hashlib
import logging
from collections import defaultdict
import numpy as np
logger = logging.getLogger(__name__)
class MemoryGraph:
"""Directed graph of observed game states.
Nodes are identified by frame hashes (MD5 of the raw 64×64 frame).
Each node stores:
- A latent vector (from the perception encoder) for similarity checks.
- The score observed at that state.
- A visit count (incremented each time the state is seen).
Edges store the action that caused the transition.
The graph resets per RESET (per level attempt), but the agent can
optionally carry summary statistics across attempts.
Attributes:
nodes: Dict mapping frame_hash → node data.
edges: Dict mapping (from_hash, to_hash) → edge data.
adjacency: Dict mapping from_hash → list of (action, to_hash).
"""
def __init__(self) -> None:
"""Initialize an empty memory graph."""
self.nodes: dict[str, dict] = {}
self.edges: dict[tuple[str, str], dict] = {}
self.adjacency: dict[str, list[tuple[str, str]]] = defaultdict(list)
@staticmethod
def hash_frame(frame: np.ndarray) -> str:
"""Compute a stable hash for a 64×64 frame.
Args:
frame: 64×64 uint8 array of color indices.
Returns:
MD5 hex string of the frame data.
"""
return hashlib.md5(frame.tobytes()).hexdigest()
def add_node(
self,
frame_hash: str,
latent: np.ndarray,
score: float = 0.0,
) -> None:
"""Add or update a node in the graph.
If the node already exists, increment its visit count.
Args:
frame_hash: Hash of the frame.
latent: Latent vector from the perception encoder.
score: Game score at this state.
"""
if frame_hash in self.nodes:
self.nodes[frame_hash]["visit_count"] += 1
self.nodes[frame_hash]["score"] = score
else:
self.nodes[frame_hash] = {
"latent": latent.copy(),
"score": score,
"visit_count": 1,
}
def add_edge(
self,
from_hash: str,
to_hash: str,
action: str,
action_data: dict | None = None,
) -> None:
"""Add a directed edge to the graph.
Args:
from_hash: Source node hash.
to_hash: Destination node hash.
action: Action name that caused this transition.
action_data: Optional action data (e.g. coordinates for ACTION6).
"""
edge_key = (from_hash, to_hash)
if edge_key in self.edges:
self.edges[edge_key]["count"] += 1
else:
self.edges[edge_key] = {
"action": action,
"action_data": action_data,
"count": 1,
}
self.adjacency[from_hash].append((action, to_hash))
def novelty(self, frame_hash: str) -> float:
"""Compute novelty score for a frame.
Novelty is high for unvisited or rarely-visited states, low for
frequently-visited ones. Used by the intrinsic reward module.
Formula: novelty = 1.0 / (1.0 + visit_count)
Args:
frame_hash: Hash of the frame to evaluate.
Returns:
Novelty score in (0, 1]. New states return 1.0.
"""
if frame_hash not in self.nodes:
return 1.0
visit_count = self.nodes[frame_hash]["visit_count"]
return 1.0 / (1.0 + visit_count)
def shortest_path_to_unexplored(self, from_hash: str) -> list[str] | None:
"""Find shortest path from a node to any unexplored action.
Uses BFS to find the nearest node that has untried actions
(actions not yet taken from that node).
Args:
from_hash: Starting node hash.
Returns:
List of action names forming the path, or None if no
unexplored action is reachable.
"""
if from_hash not in self.nodes:
return None
# BFS
from collections import deque
queue: deque[tuple[str, list[str]]] = deque([(from_hash, [])])
visited: set[str] = {from_hash}
all_actions = {"RESET", "ACTION1", "ACTION2", "ACTION3",
"ACTION4", "ACTION5", "ACTION6"}
while queue:
current_hash, path = queue.popleft()
# Check if current node has unexplored actions
tried_actions = {
edge_action
for edge_action, _ in self.adjacency.get(current_hash, [])
}
unexplored = all_actions - tried_actions
if unexplored and len(path) > 0:
return path
if unexplored:
# We're at the start node with unexplored actions
return []
# Expand neighbors
for edge_action, neighbor_hash in self.adjacency.get(current_hash, []):
if neighbor_hash not in visited:
visited.add(neighbor_hash)
queue.append((neighbor_hash, path + [edge_action]))
return None
def get_visited_states(self) -> set[str]:
"""Return the set of all visited node hashes.
Returns:
Set of frame hashes that have been visited.
"""
return set(self.nodes.keys())
def get_transition_count(self) -> int:
"""Return total number of unique transitions (edges).
Returns:
Number of edges in the graph.
"""
return len(self.edges)
def stats(self) -> dict:
"""Return summary statistics about the graph.
Returns:
Dict with node_count, edge_count, avg_visits, max_visits.
"""
visit_counts = [n["visit_count"] for n in self.nodes.values()]
return {
"node_count": len(self.nodes),
"edge_count": len(self.edges),
"avg_visits": np.mean(visit_counts) if visit_counts else 0.0,
"max_visits": max(visit_counts) if visit_counts else 0,
}
def reset(self) -> None:
"""Clear all nodes and edges.
Called when the agent RESETs a level or starts a new one.
"""
self.nodes.clear()
self.edges.clear()
self.adjacency.clear()
logger.debug("Memory graph reset")
def serialize(self) -> dict:
"""Serialize the graph for checkpointing.
Returns:
Dict representation of the graph.
"""
return {
"nodes": {
h: {"score": n["score"], "visit_count": n["visit_count"]}
for h, n in self.nodes.items()
},
"edges": {
f"{k[0]}->{k[1]}": v for k, v in self.edges.items()
},
}
|