Bridge HGF, DDM, metacognition, hierarchical to pymdp Agent

All four unique ALF modules now accept pymdp Agents directly:

- HGFPerceptualAgent: accepts pymdp Agent (converts to ALF GM internally)
- DDM bridge: add neg_efe_to_ddm() for pymdp's sign convention
- MetacognitiveAgent: accepts pymdp Agent (converts to ALF GM, wraps)
- HierarchicalGenerativeModel.from_pymdp(): creates base level from
  pymdp Agent with optional higher context levels

9 new tests cover all bridges.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: m9h

alf/ddm/__init__.py

@@ -37,6 +37,7 @@
 )
 from alf.ddm.bridge import (
     efe_to_ddm,
+    neg_efe_to_ddm,
     ddm_to_policy_probs,
 )
 from alf.ddm.fitting import (
@@ -60,6 +61,7 @@
     "simulate_ddm",
     # Bridge
     "efe_to_ddm",
+    "neg_efe_to_ddm",
     "ddm_to_policy_probs",
     # Fitting
     "DDMFitResult",

alf/ddm/bridge.py

@@ -56,6 +56,32 @@ def efe_to_ddm(
     return DDMParams(v=v, a=a, w=w, tau=jnp.array(tau))
 
 
+def neg_efe_to_ddm(
+    neg_efe: jnp.ndarray,
+    gamma: float = 4.0,
+    tau: float = 0.3,
+    base_a: float = 1.5,
+) -> DDMParams:
+    """Map pymdp neg_efe values to DDM parameters for a binary choice.
+
+    Convenience wrapper for efe_to_ddm that accepts pymdp's neg_efe
+    convention (higher = better) instead of ALF's G (lower = better).
+
+    Args:
+        neg_efe: Negative EFE for each action, shape (2,) or (1, 2).
+            Higher values indicate preferred actions.
+        gamma: Policy precision (inverse temperature).
+        tau: Non-decision time.
+        base_a: Base boundary separation.
+
+    Returns:
+        DDMParams with v, a, w, tau.
+    """
+    if neg_efe.ndim == 2:
+        neg_efe = neg_efe[0]
+    return efe_to_ddm(-neg_efe, gamma=gamma, tau=tau, base_a=base_a)
+
+
 def ddm_to_policy_probs(
     params: DDMParams,
 ) -> jnp.ndarray:

alf/hgf/bridge.py

@@ -115,6 +115,8 @@ class HGFPerceptualAgent:
 
     Args:
         gm: Generative model with B, C matrices for action selection.
+            Accepts either an ALF GenerativeModel or a pymdp Agent
+            (which will be converted to ALF GM internally).
         hgf_params: HGF parameters (BinaryHGFParams or ContinuousHGFParams).
         gamma: Policy precision (inverse temperature). Default 4.0.
         state_range: Range for discretizing HGF beliefs.
@@ -123,12 +125,15 @@ class HGFPerceptualAgent:
 
     def __init__(
         self,
-        gm: GenerativeModel,
+        gm,
         hgf_params: BinaryHGFParams | ContinuousHGFParams,
         gamma: float = 4.0,
         state_range: tuple[float, float] = (-3.0, 3.0),
         seed: int = 42,
     ):
+        if not isinstance(gm, GenerativeModel):
+            from alf.compat import pymdp_to_alf
+            gm = pymdp_to_alf(gm)
         self.gm = gm
         self.hgf_params = hgf_params
         self.gamma = gamma

alf/hierarchical.py

@@ -205,6 +205,33 @@ class HierarchicalGenerativeModel:
     def __init__(self, levels: list[HierarchicalLevel]):
         self.levels = list(levels)
 
+    @classmethod
+    def from_pymdp(cls, agent, higher_levels: list[HierarchicalLevel] = None,
+                   level_name: str = "sensorimotor") -> "HierarchicalGenerativeModel":
+        """Create a hierarchy with the lowest level from a pymdp Agent.
+
+        Extracts the A, B, C, D matrices from a pymdp Agent (stripping the
+        batch dimension) and uses them as the lowest-level (fastest) model.
+        Additional higher levels can be appended.
+
+        Args:
+            agent: pymdp.agent.Agent instance.
+            higher_levels: Optional list of HierarchicalLevel objects for
+                context/strategy levels above the base sensorimotor level.
+            level_name: Name for the base level.
+
+        Returns:
+            HierarchicalGenerativeModel with the pymdp model as level 0.
+        """
+        from alf.compat import pymdp_to_alf
+        gm = pymdp_to_alf(agent)
+        base_level = HierarchicalLevel(
+            A=gm.A[0], B=gm.B[0], C=gm.C[0], D=gm.D[0],
+            temporal_scale=1, level_name=level_name,
+        )
+        levels = [base_level] + (higher_levels or [])
+        return cls(levels)
+
     @property
     def num_levels(self) -> int:
         return len(self.levels)

alf/metacognition.py

@@ -652,10 +652,10 @@ class MetacognitiveAgent:
 
     Args:
         inner_agent: The AnalyticAgent to wrap. Can also accept a
-            GenerativeModel, in which case an AnalyticAgent is created
-            internally.
+            GenerativeModel (creates an AnalyticAgent internally) or
+            a pymdp Agent (converted to ALF GM, then wrapped).
         gamma: Initial policy precision. Only used if inner_agent is a
-            GenerativeModel. Default 4.0.
+            GenerativeModel or pymdp Agent. Default 4.0.
         monitor_decay: EMA decay for the EFEMonitor. Default 0.95.
         monitor_window: Window size for the EFEMonitor. Default 50.
         gamma_learning_rate: Learning rate for gamma adjustment based
@@ -670,7 +670,7 @@ class MetacognitiveAgent:
 
     def __init__(
         self,
-        inner_agent: AnalyticAgent | GenerativeModel,
+        inner_agent,
         gamma: float = 4.0,
         monitor_decay: float = 0.95,
         monitor_window: int = 50,
@@ -682,8 +682,13 @@ def __init__(
     ):
         if isinstance(inner_agent, GenerativeModel):
             self.agent = AnalyticAgent(inner_agent, gamma=gamma, seed=seed)
-        else:
+        elif isinstance(inner_agent, AnalyticAgent):
             self.agent = inner_agent
+        else:
+            # Assume pymdp Agent — convert to ALF
+            from alf.compat import pymdp_to_alf
+            gm = pymdp_to_alf(inner_agent)
+            self.agent = AnalyticAgent(gm, gamma=gamma, seed=seed)
 
         self.monitor = EFEMonitor(
             decay=monitor_decay,

alf/tests/test_pymdp_bridges.py

@@ -0,0 +1,137 @@
+"""Tests for pymdp bridge integration in HGF, DDM, metacognition, hierarchical."""
+
+import numpy as np
+import jax.numpy as jnp
+import pytest
+
+try:
+    from pymdp.agent import Agent as PyMDP_Agent
+    from alf.compat import alf_to_pymdp
+    HAS_PYMDP = True
+except ImportError:
+    HAS_PYMDP = False
+
+from alf.generative_model import GenerativeModel
+
+pytestmark = pytest.mark.skipif(not HAS_PYMDP, reason="pymdp not installed")
+
+
+@pytest.fixture
+def simple_model():
+    """3-state, 3-obs, 2-action model."""
+    A = np.array([[0.9, 0.05, 0.05],
+                  [0.05, 0.9, 0.05],
+                  [0.05, 0.05, 0.9]])
+    B = np.zeros((3, 3, 2))
+    B[:, :, 0] = np.array([[0.8, 0.1, 0.1],
+                            [0.1, 0.8, 0.1],
+                            [0.1, 0.1, 0.8]])
+    B[:, :, 1] = np.array([[0.1, 0.1, 0.8],
+                            [0.8, 0.1, 0.1],
+                            [0.1, 0.8, 0.1]])
+    C = np.array([2.0, 0.0, -1.0])
+    D = np.array([1/3, 1/3, 1/3])
+    return GenerativeModel(A=[A], B=[B], C=[C], D=[D])
+
+
+@pytest.fixture
+def pymdp_agent(simple_model):
+    return alf_to_pymdp(simple_model)
+
+
+class TestHGFBridgePymdp:
+    def test_accepts_pymdp_agent(self, pymdp_agent):
+        from alf.hgf.bridge import HGFPerceptualAgent
+        from alf.hgf.updates import BinaryHGFParams
+
+        params = BinaryHGFParams(omega_2=-2.0, mu_2_0=0.0, sigma_2_0=1.0)
+        agent = HGFPerceptualAgent(pymdp_agent, params, gamma=4.0)
+        assert agent.gm is not None
+        assert isinstance(agent.gm, GenerativeModel)
+
+    def test_step_works_with_pymdp(self, pymdp_agent):
+        from alf.hgf.bridge import HGFPerceptualAgent
+        from alf.hgf.updates import BinaryHGFParams
+
+        params = BinaryHGFParams(omega_2=-2.0, mu_2_0=0.0, sigma_2_0=1.0)
+        agent = HGFPerceptualAgent(pymdp_agent, params, gamma=4.0)
+        action, info = agent.step(1.0)
+        assert isinstance(action, int)
+        assert "G" in info
+
+
+class TestDDMBridgePymdp:
+    def test_neg_efe_to_ddm(self):
+        from alf.ddm.bridge import neg_efe_to_ddm, efe_to_ddm
+
+        G = jnp.array([-1.5, -0.5])
+        neg_efe = jnp.array([1.5, 0.5])
+
+        params_from_G = efe_to_ddm(G)
+        params_from_neg_efe = neg_efe_to_ddm(neg_efe)
+
+        np.testing.assert_allclose(float(params_from_G.v),
+                                   float(params_from_neg_efe.v), atol=1e-5)
+        np.testing.assert_allclose(float(params_from_G.a),
+                                   float(params_from_neg_efe.a), atol=1e-5)
+
+    def test_neg_efe_to_ddm_with_batch_dim(self):
+        from alf.ddm.bridge import neg_efe_to_ddm
+
+        neg_efe_batched = jnp.array([[1.5, 0.5]])
+        params = neg_efe_to_ddm(neg_efe_batched)
+        assert jnp.isfinite(params.v)
+
+
+class TestMetacognitionPymdp:
+    def test_accepts_pymdp_agent(self, pymdp_agent):
+        from alf.metacognition import MetacognitiveAgent
+
+        agent = MetacognitiveAgent(pymdp_agent, gamma=4.0)
+        assert agent.gm is not None
+
+    def test_step_works_with_pymdp(self, pymdp_agent):
+        from alf.metacognition import MetacognitiveAgent
+
+        agent = MetacognitiveAgent(pymdp_agent, gamma=4.0)
+        action, info = agent.step([0])
+        assert isinstance(action, int)
+        assert "metacognitive_confidence" in info
+        assert 0.0 <= info["metacognitive_confidence"] <= 1.0
+
+    def test_learn_works_with_pymdp(self, pymdp_agent):
+        from alf.metacognition import MetacognitiveAgent
+
+        agent = MetacognitiveAgent(pymdp_agent, gamma=4.0)
+        agent.step([0])
+        agent.learn(1.0)
+        assert len(agent.accuracy_history) == 1
+
+
+class TestHierarchicalPymdp:
+    def test_from_pymdp(self, pymdp_agent):
+        from alf.hierarchical import HierarchicalGenerativeModel
+
+        hierarchy = HierarchicalGenerativeModel.from_pymdp(pymdp_agent)
+        assert hierarchy.num_levels == 1
+        assert hierarchy.levels[0].num_states == 3
+        assert hierarchy.levels[0].num_actions == 2
+
+    def test_from_pymdp_with_higher_levels(self, pymdp_agent):
+        from alf.hierarchical import HierarchicalGenerativeModel, HierarchicalLevel
+
+        context_A = np.eye(2)
+        context_B = np.stack([np.eye(2)] * 2, axis=-1)
+        context_C = np.array([1.0, 0.0])
+        context_D = np.array([0.5, 0.5])
+        context_level = HierarchicalLevel(
+            A=context_A, B=context_B, C=context_C, D=context_D,
+            temporal_scale=5, level_name="context",
+        )
+
+        hierarchy = HierarchicalGenerativeModel.from_pymdp(
+            pymdp_agent, higher_levels=[context_level]
+        )
+        assert hierarchy.num_levels == 2
+        assert hierarchy.levels[0].level_name == "sensorimotor"
+        assert hierarchy.levels[1].level_name == "context"