Skip to content

Baselines API Reference

Generated directly from docstrings and source. See concepts/marl.md and specs/algorithm-spec.md for the conceptual/contract writeups.

BaseAlgorithm

base.BaseAlgorithm

Bases: ABC

Stable learning interface. Algorithms treat env as a black box.

Source code in src/ppage/baselines/base.py 
 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
class BaseAlgorithm(ABC):
    """
    Stable learning interface.
    Algorithms treat env as a black box.
    """

    def __init__(self, env, config: dict):
        self.env = env
        self.config = config

    @abstractmethod
    def select_actions(self, observations: dict) -> dict:
        """
        observations: {agent_name: obs_dict}
        returns: {agent_name: action_int}
        """

    @abstractmethod
    def train(self):
        pass

    def evaluate(self, episodes: int = 5, max_steps: int = 500) -> dict:
        """Run greedy-ish evaluation episodes and return summary metrics.

        Returns a dict with the number of episodes, mean episode length, and
        mean per-agent return. Callers (the run_* scripts) print this; the
        method no longer discards its results.
        """
        episode_lengths = []
        agent_returns = defaultdict(list)

        for _ in range(episodes):
            obs, _ = self.env.reset()
            done = False
            steps = 0
            ep_reward = defaultdict(float)
            while not done and steps < max_steps:
                actions = self.select_actions(obs)
                step_out = self.env.step(actions)
                obs = step_out["obs"]
                for agent_id, r in step_out["reward"].items():
                    ep_reward[agent_id] += float(r)
                done = step_out["terminated"] or step_out["truncated"]
                steps += 1
            episode_lengths.append(steps)
            for agent_id, total in ep_reward.items():
                agent_returns[agent_id].append(total)

        summary = {
            "episodes": episodes,
            "mean_episode_length": (
                float(np.mean(episode_lengths)) if episode_lengths else 0.0
            ),
        }
        for agent_id, returns in agent_returns.items():
            summary[f"mean_return_{agent_id}"] = float(np.mean(returns))
        return summary

select_actions abstractmethod 

select_actions(observations: dict) -> dict

observations: {agent_name: obs_dict} returns: {agent_name: action_int}

Source code in src/ppage/baselines/base.py 
19
20
21
22
23
24
@abstractmethod
def select_actions(self, observations: dict) -> dict:
    """
    observations: {agent_name: obs_dict}
    returns: {agent_name: action_int}
    """

evaluate 

evaluate(episodes: int = 5, max_steps: int = 500) -> dict

Run greedy-ish evaluation episodes and return summary metrics.

Returns a dict with the number of episodes, mean episode length, and mean per-agent return. Callers (the run_* scripts) print this; the method no longer discards its results.

Source code in src/ppage/baselines/base.py 
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
def evaluate(self, episodes: int = 5, max_steps: int = 500) -> dict:
    """Run greedy-ish evaluation episodes and return summary metrics.

    Returns a dict with the number of episodes, mean episode length, and
    mean per-agent return. Callers (the run_* scripts) print this; the
    method no longer discards its results.
    """
    episode_lengths = []
    agent_returns = defaultdict(list)

    for _ in range(episodes):
        obs, _ = self.env.reset()
        done = False
        steps = 0
        ep_reward = defaultdict(float)
        while not done and steps < max_steps:
            actions = self.select_actions(obs)
            step_out = self.env.step(actions)
            obs = step_out["obs"]
            for agent_id, r in step_out["reward"].items():
                ep_reward[agent_id] += float(r)
            done = step_out["terminated"] or step_out["truncated"]
            steps += 1
        episode_lengths.append(steps)
        for agent_id, total in ep_reward.items():
            agent_returns[agent_id].append(total)

    summary = {
        "episodes": episodes,
        "mean_episode_length": (
            float(np.mean(episode_lengths)) if episode_lengths else 0.0
        ),
    }
    for agent_id, returns in agent_returns.items():
        summary[f"mean_return_{agent_id}"] = float(np.mean(returns))
    return summary

IQL

iql.IQL

Bases: BaseAlgorithm

Tabular Independent Q-Learning.

One Q-table per agent. State is encoded as a hashable tuple from the raw observation dict — no assumption about observation structure, fully wrapper-compatible.

Source code in src/ppage/baselines/IQL/iql.py 
 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
class IQL(BaseAlgorithm):
    """
    Tabular Independent Q-Learning.

    One Q-table per agent. State is encoded as a hashable tuple from the raw
    observation dict — no assumption about observation structure, fully
    wrapper-compatible.
    """

    def __init__(self, env, config):
        super().__init__(env, config)

        self.alpha = config.get("alpha", 0.1)
        self.gamma = config.get("gamma", 0.99)
        self.epsilon = config.get("epsilon", 0.1)
        self.episodes = config.get("episodes", 500)
        self.epsilon_decay = config.get("epsilon_decay", 1.0)
        self.min_epsilon = config.get("min_epsilon", 0.01)
        self.action_dim = config.get("action_dim", 5)

        # seeded RNG for reproducible exploration
        self.rng = default_rng(config.get("seed", None))

        # discover agent IDs from an initial reset
        initial_obs, _ = self.env.reset()
        self.agent_ids = list(initial_obs.keys())

        # per-agent Q-tables backed by defaultdict (supports arbitrary states)
        self.q_tables = {
            agent_id: defaultdict(lambda: np.zeros(self.action_dim))
            for agent_id in self.agent_ids
        }

    # ------------------------------------------------------------------
    # State encoding
    # ------------------------------------------------------------------

    def _encode_state(self, obs_dict) -> tuple:
        """Convert observation dict into a hashable tuple (wrapper-safe)."""

        def _convert(obj):
            if isinstance(obj, dict):
                return tuple(sorted((k, _convert(v)) for k, v in obj.items()))
            if isinstance(obj, (list, tuple)):
                return tuple(_convert(v) for v in obj)
            if hasattr(obj, "tolist"):
                val = obj.tolist()
                # numpy array → list; numpy scalar → Python int/float
                if isinstance(val, (list, tuple)):
                    return tuple(_convert(v) for v in val)
                return val
            return obj

        return _convert(obs_dict)

    # ------------------------------------------------------------------
    # Action selection (BaseAlgorithm interface)
    # ------------------------------------------------------------------

    def select_actions(self, observations: dict) -> dict:
        actions = {}
        for agent_id, obs in observations.items():
            state = self._encode_state(obs)
            if self.rng.random() < self.epsilon:
                action = int(self.rng.integers(self.action_dim))
            else:
                action = int(np.argmax(self.q_tables[agent_id][state]))
            actions[agent_id] = action
        return actions

    # ------------------------------------------------------------------
    # Training loop (BaseAlgorithm interface)
    # ------------------------------------------------------------------

    def train(self):
        for ep in range(self.episodes):
            obs, _ = self.env.reset()
            done = False
            ep_reward = {agent_id: 0.0 for agent_id in self.agent_ids}

            while not done:
                actions = self.select_actions(obs)

                step_out = self.env.step(actions)
                next_obs = step_out["obs"]
                rewards = step_out["reward"]
                done = step_out["terminated"] or step_out["truncated"]
                # Only a true terminal state cuts the bootstrap; a timeout
                # truncation does not (the episode is cut for time, not because
                # the world ended), so keep bootstrapping through it.
                terminal = bool(step_out["terminated"])

                for agent_id in self.agent_ids:
                    s = self._encode_state(obs[agent_id])
                    a = actions[agent_id]
                    r = float(rewards[agent_id])
                    ep_reward[agent_id] += r
                    s_next = self._encode_state(next_obs[agent_id])

                    q_current = self.q_tables[agent_id][s][a]
                    q_next_max = (
                        0.0
                        if terminal
                        else float(np.max(self.q_tables[agent_id][s_next]))
                    )
                    td_error = r + self.gamma * q_next_max - q_current
                    self.q_tables[agent_id][s][a] += self.alpha * td_error

                obs = next_obs

            self.epsilon = max(self.min_epsilon, self.epsilon * self.epsilon_decay)

            if (ep + 1) % 100 == 0:
                reward_str = ", ".join(f"{k}={v:.2f}" for k, v in ep_reward.items())
                LOGGER.info(
                    "Episode %d/%d | eps=%.3f | %s",
                    ep + 1,
                    self.episodes,
                    self.epsilon,
                    reward_str,
                )

    # ------------------------------------------------------------------
    # Persistence
    # ------------------------------------------------------------------

    def save(self, path: str) -> None:
        """Serialize Q-tables to a pickle file."""
        os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
        payload = {agent_id: dict(table) for agent_id, table in self.q_tables.items()}
        with open(path, "wb") as fh:
            pickle.dump(payload, fh)
        LOGGER.info("Saved Q-tables -> %s", path)

    @classmethod
    def load(cls, env, config: dict, path: str) -> "IQL":
        """Load Q-tables from a pickle file; returns a ready-to-evaluate instance."""
        instance = cls(env, config)
        with open(path, "rb") as fh:
            payload: dict = pickle.load(fh)
        for agent_id, table in payload.items():
            instance.q_tables[agent_id].update(table)
        LOGGER.info("Loaded Q-tables from %s (agents: %s)", path, list(payload.keys()))
        return instance

save 

save(path: str) -> None

Serialize Q-tables to a pickle file.

Source code in src/ppage/baselines/IQL/iql.py 
143
144
145
146
147
148
149
def save(self, path: str) -> None:
    """Serialize Q-tables to a pickle file."""
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    payload = {agent_id: dict(table) for agent_id, table in self.q_tables.items()}
    with open(path, "wb") as fh:
        pickle.dump(payload, fh)
    LOGGER.info("Saved Q-tables -> %s", path)

load classmethod 

load(env, config: dict, path: str) -> IQL

Load Q-tables from a pickle file; returns a ready-to-evaluate instance.

Source code in src/ppage/baselines/IQL/iql.py 
151
152
153
154
155
156
157
158
159
160
@classmethod
def load(cls, env, config: dict, path: str) -> "IQL":
    """Load Q-tables from a pickle file; returns a ready-to-evaluate instance."""
    instance = cls(env, config)
    with open(path, "rb") as fh:
        payload: dict = pickle.load(fh)
    for agent_id, table in payload.items():
        instance.q_tables[agent_id].update(table)
    LOGGER.info("Loaded Q-tables from %s (agents: %s)", path, list(payload.keys()))
    return instance

CQL

cql.CQL

Bases: BaseAlgorithm

Tabular Centralized Q-Learning.

Single shared Q-table: joint_state → Q-vector of length action_dim^n_agents. Joint state : tuple of every agent's encoded observation (ordered by agent_ids). Joint action : integer encoding all agents' actions in base action_dim. Central reward: sum of all agents' rewards for each transition.

Source code in src/ppage/baselines/CQL/cql.py 
 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
class CQL(BaseAlgorithm):
    """
    Tabular Centralized Q-Learning.

    Single shared Q-table: joint_state → Q-vector of length action_dim^n_agents.
    Joint state  : tuple of every agent's encoded observation (ordered by agent_ids).
    Joint action : integer encoding all agents' actions in base action_dim.
    Central reward: sum of all agents' rewards for each transition.
    """

    def __init__(self, env, config: dict):
        super().__init__(env, config)

        self.alpha = config.get("alpha", 0.1)
        self.gamma = config.get("gamma", 0.99)
        self.epsilon = config.get("epsilon", 1.0)
        self.epsilon_decay = config.get("epsilon_decay", 0.995)
        self.min_epsilon = config.get("min_epsilon", 0.05)
        self.action_dim = config.get("action_dim", 5)
        self.episodes = config.get("episodes", 1000)

        self.rng = default_rng(config.get("seed", None))

        initial_obs, _ = self.env.reset()
        self.agent_ids = list(initial_obs.keys())
        self.n_agents = len(self.agent_ids)
        self.n_joint_actions = self.action_dim**self.n_agents

        # shape for reshaping the Q-vector into a per-agent tensor
        self._action_shape = (self.action_dim,) * self.n_agents

        # single shared Q-table: joint_state (tuple) → np.zeros(n_joint_actions)
        self.q_table: defaultdict = defaultdict(
            lambda: np.zeros(self.n_joint_actions, dtype=np.float32)
        )

    # ------------------------------------------------------------------
    # State encoding (wrapper-compatible, same approach as IQL)
    # ------------------------------------------------------------------

    def _encode_state(self, obs_dict) -> tuple:
        def _convert(obj):
            if isinstance(obj, dict):
                return tuple(sorted((k, _convert(v)) for k, v in obj.items()))
            if isinstance(obj, (list, tuple)):
                return tuple(_convert(v) for v in obj)
            if hasattr(obj, "tolist"):
                val = obj.tolist()
                if isinstance(val, (list, tuple)):
                    return tuple(_convert(v) for v in val)
                return val
            return obj

        return _convert(obs_dict)

    def _joint_state(self, observations: dict) -> tuple:
        """Encode all agents' observations into a single hashable joint state."""
        return tuple(self._encode_state(observations[aid]) for aid in self.agent_ids)

    # ------------------------------------------------------------------
    # Joint action encoding / decoding
    # ------------------------------------------------------------------

    def _joint_action_index(self, actions: dict) -> int:
        """Convert per-agent action dict → single joint-action integer."""
        idx = 0
        for aid in self.agent_ids:
            idx = idx * self.action_dim + int(actions[aid])
        return idx

    # ------------------------------------------------------------------
    # Action selection
    # ------------------------------------------------------------------

    def select_actions(self, observations: dict) -> dict:
        joint_s = self._joint_state(observations)
        q_vals = self.q_table[joint_s]  # (n_joint_actions,)
        q_tensor = q_vals.reshape(self._action_shape)  # (a_dim, a_dim, ...)

        actions = {}
        for i, aid in enumerate(self.agent_ids):
            if self.rng.random() < self.epsilon:
                actions[aid] = int(self.rng.integers(self.action_dim))
            else:
                # marginalise over all other agents' action axes
                axes = tuple(j for j in range(self.n_agents) if j != i)
                q_marginal = q_tensor.mean(axis=axes) if axes else q_vals
                actions[aid] = int(np.argmax(q_marginal))
        return actions

    # ------------------------------------------------------------------
    # Training loop
    # ------------------------------------------------------------------

    def train(self):
        for ep in range(self.episodes):
            obs, _ = self.env.reset()
            done = False
            ep_reward = {aid: 0.0 for aid in self.agent_ids}

            while not done:
                actions = self.select_actions(obs)
                step_out = self.env.step(actions)
                next_obs = step_out["obs"]
                rewards = step_out["reward"]
                done = step_out["terminated"] or step_out["truncated"]
                # Only a true terminal state cuts the bootstrap; a timeout
                # truncation does not.
                terminal = bool(step_out["terminated"])

                joint_s = self._joint_state(obs)
                joint_s_next = self._joint_state(next_obs)
                joint_a = self._joint_action_index(actions)

                # central reward = sum of all agents' rewards
                central_r = 0.0
                for aid in self.agent_ids:
                    r = float(rewards[aid])
                    ep_reward[aid] += r
                    central_r += r

                # centralized TD update on the shared Q-table
                q_current = self.q_table[joint_s][joint_a]
                q_next_max = (
                    0.0 if terminal else float(np.max(self.q_table[joint_s_next]))
                )
                td_error = central_r + self.gamma * q_next_max - q_current
                self.q_table[joint_s][joint_a] += self.alpha * td_error

                obs = next_obs

            self.epsilon = max(self.min_epsilon, self.epsilon * self.epsilon_decay)

            if (ep + 1) % 100 == 0:
                reward_str = ", ".join(f"{k}={v:.2f}" for k, v in ep_reward.items())
                LOGGER.info(
                    "Episode %d/%d | eps=%.3f | joint_states=%d | %s",
                    ep + 1,
                    self.episodes,
                    self.epsilon,
                    len(self.q_table),
                    reward_str,
                )

    # ------------------------------------------------------------------
    # Persistence
    # ------------------------------------------------------------------

    def save(self, path: str) -> None:
        os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
        payload = {
            "q_table": dict(self.q_table),
            "agent_ids": self.agent_ids,
            "n_joint_actions": self.n_joint_actions,
        }
        with open(path, "wb") as fh:
            pickle.dump(payload, fh)
        LOGGER.info("Saved -> %s  (%d joint states)", path, len(self.q_table))

    @classmethod
    def load(cls, env, config: dict, path: str) -> "CQL":
        instance = cls(env, config)
        with open(path, "rb") as fh:
            payload: dict = pickle.load(fh)
        instance.q_table.update(payload["q_table"])
        LOGGER.info("Loaded from %s  (%d joint states)", path, len(instance.q_table))
        return instance

MixedTrainer

mix_train.MixedTrainer

Bases: BaseAlgorithm

Per-team algorithm assignment: predators and prey can use different algorithms.

predator_algo / prey_algo: "iql" or "cql" - iql : independent — one Q-table per agent in the team - cql : centralized — one shared Q-table over joint team state-action space; action selection marginalises the joint Q over other team members

Source code in src/ppage/baselines/MIXED/mix_train.py 
 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
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
class MixedTrainer(BaseAlgorithm):
    """
    Per-team algorithm assignment: predators and prey can use different algorithms.

    predator_algo / prey_algo: "iql" or "cql"
    - iql : independent — one Q-table per agent in the team
    - cql : centralized — one shared Q-table over joint team state-action space;
            action selection marginalises the joint Q over other team members
    """

    def __init__(self, env, config: dict):
        super().__init__(env, config)

        self.alpha = config.get("alpha", 0.1)
        self.gamma = config.get("gamma", 0.99)
        self.epsilon = config.get("epsilon", 1.0)
        self.epsilon_decay = config.get("epsilon_decay", 0.995)
        self.min_epsilon = config.get("min_epsilon", 0.05)
        self.action_dim = config.get("action_dim", 5)
        self.episodes = config.get("episodes", 1000)
        self.predator_algo = config.get("predator_algo", "cql").lower()
        self.prey_algo = config.get("prey_algo", "iql").lower()

        self.rng = default_rng(config.get("seed", None))

        initial_obs, _ = self.env.reset()
        self.agent_ids = list(initial_obs.keys())

        # Partition agents into teams
        self._predators = [
            ag.agent_name
            for ag in self.env.agents
            if ag.agent_type.startswith(_PREDATOR_TEAM)
        ]
        self._prey = [
            ag.agent_name
            for ag in self.env.agents
            if not ag.agent_type.startswith(_PREDATOR_TEAM)
        ]

        # Build team lookup: agent_name → team_key
        self._team_of: dict[str, str] = {}
        for aid in self._predators:
            self._team_of[aid] = _PREDATOR_TEAM
        for aid in self._prey:
            self._team_of[aid] = _PREY_TEAM

        self._algo_of: dict[str, str] = {
            _PREDATOR_TEAM: self.predator_algo,
            _PREY_TEAM: self.prey_algo,
        }

        # ------------------------------------------------------------------
        # Q-table structures
        # ------------------------------------------------------------------

        # IQL tables: one per agent
        self._iql_tables: dict[str, defaultdict] = {}

        # CQL tables: one per team
        self._cql_tables: dict[str, defaultdict] = {}
        self._cql_team_ids: dict[str, list[str]] = {}  # team → ordered agent list
        self._cql_n_joint: dict[str, int] = {}  # team → n_joint_actions
        self._cql_action_shape: dict[str, tuple] = {}  # team → reshape tuple

        for team_key, members, algo in [
            (_PREDATOR_TEAM, self._predators, self.predator_algo),
            (_PREY_TEAM, self._prey, self.prey_algo),
        ]:
            if algo == "iql":
                for aid in members:
                    self._iql_tables[aid] = defaultdict(
                        lambda: np.zeros(self.action_dim)
                    )
            else:  # cql
                n_joint = self.action_dim ** len(members)
                self._cql_team_ids[team_key] = list(members)
                self._cql_n_joint[team_key] = n_joint
                self._cql_action_shape[team_key] = (self.action_dim,) * len(members)
                self._cql_tables[team_key] = defaultdict(
                    lambda nj=n_joint: np.zeros(nj, dtype=np.float32)
                )

    # ------------------------------------------------------------------
    # State encoding (wrapper-safe, identical to IQL)
    # ------------------------------------------------------------------

    def _encode_state(self, obs_dict) -> tuple:
        def _convert(obj):
            if isinstance(obj, dict):
                return tuple(sorted((k, _convert(v)) for k, v in obj.items()))
            if isinstance(obj, (list, tuple)):
                return tuple(_convert(v) for v in obj)
            if hasattr(obj, "tolist"):
                val = obj.tolist()
                if isinstance(val, (list, tuple)):
                    return tuple(_convert(v) for v in val)
                return val
            return obj

        return _convert(obs_dict)

    def _team_joint_state(self, observations: dict, team_key: str) -> tuple:
        return tuple(
            self._encode_state(observations[aid])
            for aid in self._cql_team_ids[team_key]
        )

    def _team_joint_action_index(self, actions: dict, team_key: str) -> int:
        idx = 0
        for aid in self._cql_team_ids[team_key]:
            idx = idx * self.action_dim + int(actions[aid])
        return idx

    # ------------------------------------------------------------------
    # Action selection
    # ------------------------------------------------------------------

    def select_actions(self, observations: dict) -> dict:
        actions = {}

        for team_key, members, algo in [
            (_PREDATOR_TEAM, self._predators, self.predator_algo),
            (_PREY_TEAM, self._prey, self.prey_algo),
        ]:
            if algo == "iql":
                for aid in members:
                    s = self._encode_state(observations[aid])
                    if self.rng.random() < self.epsilon:
                        actions[aid] = int(self.rng.integers(self.action_dim))
                    else:
                        actions[aid] = int(np.argmax(self._iql_tables[aid][s]))
            else:  # cql — marginalise joint Q
                joint_s = self._team_joint_state(observations, team_key)
                q_vals = self._cql_tables[team_key][joint_s]
                q_tensor = q_vals.reshape(self._cql_action_shape[team_key])
                team_members = self._cql_team_ids[team_key]
                n = len(team_members)

                for i, aid in enumerate(team_members):
                    if self.rng.random() < self.epsilon:
                        actions[aid] = int(self.rng.integers(self.action_dim))
                    else:
                        axes = tuple(j for j in range(n) if j != i)
                        q_marginal = q_tensor.mean(axis=axes) if axes else q_vals
                        actions[aid] = int(np.argmax(q_marginal))

        return actions

    # ------------------------------------------------------------------
    # Training loop
    # ------------------------------------------------------------------

    def train(self):
        for ep in range(self.episodes):
            obs, _ = self.env.reset()
            done = False
            ep_reward = {aid: 0.0 for aid in self.agent_ids}

            while not done:
                actions = self.select_actions(obs)
                step_out = self.env.step(actions)
                next_obs = step_out["obs"]
                rewards = step_out["reward"]
                done = step_out["terminated"] or step_out["truncated"]
                # Only a true terminal state cuts the bootstrap; a timeout
                # truncation does not.
                terminal = bool(step_out["terminated"])

                for aid in self.agent_ids:
                    ep_reward[aid] += float(rewards[aid])

                # IQL updates — per agent
                for aid, table in self._iql_tables.items():
                    s = self._encode_state(obs[aid])
                    s_next = self._encode_state(next_obs[aid])
                    a = actions[aid]
                    r = float(rewards[aid])
                    q_current = table[s][a]
                    q_next_max = 0.0 if terminal else float(np.max(table[s_next]))
                    td_error = r + self.gamma * q_next_max - q_current
                    table[s][a] += self.alpha * td_error

                # CQL updates — per team (centralized)
                for team_key, table in self._cql_tables.items():
                    joint_s = self._team_joint_state(obs, team_key)
                    joint_s_next = self._team_joint_state(next_obs, team_key)
                    joint_a = self._team_joint_action_index(actions, team_key)

                    # central reward = sum of team members' rewards
                    central_r = sum(
                        float(rewards[aid]) for aid in self._cql_team_ids[team_key]
                    )

                    q_current = table[joint_s][joint_a]
                    q_next_max = 0.0 if terminal else float(np.max(table[joint_s_next]))
                    td_error = central_r + self.gamma * q_next_max - q_current
                    table[joint_s][joint_a] += self.alpha * td_error

                obs = next_obs

            self.epsilon = max(self.min_epsilon, self.epsilon * self.epsilon_decay)

            if (ep + 1) % 100 == 0:
                reward_str = ", ".join(f"{k}={v:.2f}" for k, v in ep_reward.items())
                LOGGER.info(
                    "Episode %d/%d | eps=%.3f | %s",
                    ep + 1,
                    self.episodes,
                    self.epsilon,
                    reward_str,
                )

    # ------------------------------------------------------------------
    # Persistence
    # ------------------------------------------------------------------

    def save(self, path: str) -> None:
        os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
        payload = {
            "iql_tables": {aid: dict(t) for aid, t in self._iql_tables.items()},
            "cql_tables": {tk: dict(t) for tk, t in self._cql_tables.items()},
            "cql_team_ids": self._cql_team_ids,
            "algo_of": self._algo_of,
        }
        with open(path, "wb") as fh:
            pickle.dump(payload, fh)
        LOGGER.info("Saved -> %s", path)

    @classmethod
    def load(cls, env, config: dict, path: str) -> "MixedTrainer":
        instance = cls(env, config)
        with open(path, "rb") as fh:
            payload: dict = pickle.load(fh)
        for aid, table in payload["iql_tables"].items():
            if aid in instance._iql_tables:
                instance._iql_tables[aid].update(table)
        for tk, table in payload["cql_tables"].items():
            if tk in instance._cql_tables:
                instance._cql_tables[tk].update(table)
        LOGGER.info("Loaded from %s", path)
        return instance

DQN

dqn.DQN

Bases: BaseAlgorithm

Independent DQN: one Q-network, target network, and replay buffer per agent.

Source code in src/ppage/baselines/DQN/dqn.py 
 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
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
class DQN(BaseAlgorithm):
    """Independent DQN: one Q-network, target network, and replay buffer per agent."""

    def __init__(self, env, config: dict):
        super().__init__(env, config)

        self.gamma = float(config.get("gamma", 0.99))
        self.epsilon = float(config.get("epsilon", 1.0))
        self.epsilon_decay = float(config.get("epsilon_decay", 0.995))
        self.min_epsilon = float(config.get("min_epsilon", 0.05))
        self.episodes = int(config.get("episodes", 1000))
        self.batch_size = int(config.get("batch_size", 32))
        self.buffer_size = int(config.get("buffer_size", 10000))
        self.min_replay_size = int(config.get("min_replay_size", self.batch_size))
        self.target_update_interval = int(config.get("target_update_interval", 100))
        self.learning_rate = float(config.get("learning_rate", 1e-3))
        self.hidden_layers = [int(v) for v in config.get("hidden_layers", [128, 128])]
        self.grad_clip = float(config.get("grad_clip", 5.0))
        self.device = torch.device(config.get("device", "cpu"))
        self.verbose = bool(config.get("verbose", True))
        self.log_interval = int(config.get("log_interval", 10))
        self.debug_first_episode = bool(config.get("debug_first_episode", True))
        self.save_path = config.get("save_path", None)
        self.curves_path: Optional[str] = config.get("curves_path", None)
        self.double_dqn = bool(config.get("double_dqn", False))  # ADDED
        self.dueling = bool(config.get("dueling", False))  # ADDED

        seed = config.get("seed", None)
        self.rng = default_rng(seed)
        if seed is not None:
            torch.manual_seed(int(seed))

        # -- observation encoder contract (attached by build_environment) --
        self.observation_encoder = getattr(self.env, "observation_encoder", None)
        if self.observation_encoder is None:
            raise ValueError(
                "Environment is missing observation_encoder. Attach an observation "
                "builder's encode() method to env.observation_encoder before "
                "constructing DQN (see run_from_config.build_environment)."
            )

        initial_obs, _ = self.env.reset()
        self.agent_ids = list(initial_obs.keys())
        self.state_dim = self._encode_observation(initial_obs[self.agent_ids[0]]).shape[
            0
        ]
        self.action_dim = self._resolve_action_dim(config)

        if self.buffer_size < self.batch_size:
            warnings.warn(
                f"buffer_size ({self.buffer_size}) is smaller than batch_size "
                f"({self.batch_size}) -- training will never fill a full batch.",
                stacklevel=2,
            )

        self._build_learners(seed)
        self._train_steps = 0

        self._debug(
            "Initialized DQN | "
            f"agents={self.agent_ids} | state_dim={self.state_dim} | "
            f"action_dim={self.action_dim} | device={self.device} | "
            f"batch_size={self.batch_size} | min_replay_size={self.min_replay_size}"
        )

    # ------------------------------------------------------------------
    # Setup helpers
    # ------------------------------------------------------------------

    def _resolve_action_dim(self, config: dict) -> int:
        """Infer action_dim from the env's action plugin, or validate if configured."""
        plugin = getattr(self.env, "action_space_plugin", None)
        plugin_n_actions = (
            int(plugin.n_actions)
            if plugin is not None
            else int(self.env.action_space.n)
        )

        configured = config.get("action_dim", None)
        if configured is None:
            return plugin_n_actions

        configured = int(configured)
        if configured != plugin_n_actions:
            raise ValueError(
                f"DQN config error: 'action_dim'={configured} does not match the "
                f"environment's action space size ({plugin_n_actions}). Fix "
                "'action_dim' in your experiment YAML, or remove it entirely so "
                "it gets inferred automatically."
            )
        return configured

    def _build_learners(self, seed) -> None:
        self.q_networks = {}
        self.target_networks = {}
        self.optimizers = {}
        self.replay_buffers = {}

        network_cls = (
            DuelingQNetwork if self.dueling else QNetwork
        )  # flag to enable duel DQN

        for i, agent_id in enumerate(self.agent_ids):
            buffer_seed = None if seed is None else int(seed) + i

            self.q_networks[agent_id] = network_cls(
                self.state_dim, self.hidden_layers, self.action_dim
            ).to(self.device)
            self.target_networks[agent_id] = network_cls(
                self.state_dim, self.hidden_layers, self.action_dim
            ).to(self.device)
            self.target_networks[agent_id].load_state_dict(
                self.q_networks[agent_id].state_dict()
            )
            self.target_networks[agent_id].eval()

            self.optimizers[agent_id] = optim.Adam(
                self.q_networks[agent_id].parameters(), lr=self.learning_rate
            )
            self.replay_buffers[agent_id] = ReplayBuffer(
                self.buffer_size, self.state_dim, seed=buffer_seed
            )

    def _debug(self, message: str) -> None:
        if self.verbose:
            print(f"[DQN] {message}")

    # ------------------------------------------------------------------
    # Observation encoding
    # ------------------------------------------------------------------

    def _encode_observation(self, obs) -> np.ndarray:
        encoded = self.observation_encoder(obs, self.env)
        return np.asarray(encoded, dtype=np.float32).reshape(-1)

    def _validate_state_shape(self, agent_id: str, state: np.ndarray) -> None:
        if state.shape[0] != self.state_dim:
            raise ValueError(
                f"DQN expected state_dim={self.state_dim}, but agent {agent_id!r} "
                f"produced state_dim={state.shape[0]}. This usually means the "
                "observation plugin is returning variable-length observations."
            )

    # ------------------------------------------------------------------
    # Action selection (BaseAlgorithm interface)
    # ------------------------------------------------------------------

    def select_actions(self, observations: dict) -> dict:
        actions = {}
        for agent_id, obs in observations.items():
            if self.rng.random() < self.epsilon:
                actions[agent_id] = int(self.rng.integers(self.action_dim))
                continue

            state = self._encode_observation(obs)
            state_t = torch.from_numpy(state).float().unsqueeze(0).to(self.device)
            with torch.no_grad():
                q_values = self.q_networks[agent_id](state_t)
            actions[agent_id] = int(torch.argmax(q_values, dim=1).item())
        return actions

    # ------------------------------------------------------------------
    # Learning
    # ------------------------------------------------------------------

    def _optimize_agent(self, agent_id: str) -> float | None:
        buffer = self.replay_buffers[agent_id]
        if len(buffer) < max(self.min_replay_size, self.batch_size):
            return None

        states, actions, rewards, next_states, dones = buffer.sample(self.batch_size)
        states_t = torch.from_numpy(states).to(self.device)
        actions_t = torch.from_numpy(actions).long().unsqueeze(1).to(self.device)
        rewards_t = torch.from_numpy(rewards).to(self.device)
        next_states_t = torch.from_numpy(next_states).to(self.device)
        dones_t = torch.from_numpy(dones.astype(np.float32)).to(self.device)

        q_values = self.q_networks[agent_id](states_t).gather(1, actions_t).squeeze(1)
        with torch.no_grad():
            if self.double_dqn:
                # action selection from the ONLINE network, evaluation from the TARGET
                # network: fixes vanilla DQN's overestimation bias.
                next_actions = self.q_networks[agent_id](next_states_t).argmax(
                    dim=1, keepdim=True
                )
                next_q_values = (
                    self.target_networks[agent_id](next_states_t)
                    .gather(1, next_actions)
                    .squeeze(1)
                )
            else:
                next_q_values = self.target_networks[agent_id](next_states_t).max(
                    dim=1
                )[0]

            targets = rewards_t + self.gamma * (1.0 - dones_t) * next_q_values

        loss = nn.SmoothL1Loss()(q_values, targets)

        optimizer = self.optimizers[agent_id]
        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(
            self.q_networks[agent_id].parameters(), self.grad_clip
        )
        optimizer.step()

        # counts total gradient steps across all agents — with N agents this
        # triggers a sync every target_update_interval/N env steps per agent.
        self._train_steps += 1
        if self._train_steps % self.target_update_interval == 0:
            self._sync_targets()

        return float(loss.item())

    def _sync_targets(self) -> None:
        for agent_id in self.agent_ids:
            self.target_networks[agent_id].load_state_dict(
                self.q_networks[agent_id].state_dict()
            )

    # ------------------------------------------------------------------
    # Training loop (BaseAlgorithm interface)
    # ------------------------------------------------------------------

    def train(self):
        recorder = (
            CurveRecorder(self.curves_path, self.agent_ids)
            if self.curves_path
            else None
        )

        self._debug(f"Starting training for {self.episodes} episodes")
        try:
            self._train_loop(recorder)
        finally:
            if recorder:
                recorder.close()

        if self.save_path:
            self.save(self.save_path)

    def _train_loop(self, recorder: Optional[CurveRecorder]) -> None:
        for episode in range(self.episodes):
            observations, _ = self.env.reset()
            episode_rewards = {agent_id: 0.0 for agent_id in self.agent_ids}
            episode_losses = {agent_id: [] for agent_id in self.agent_ids}
            done = False
            step_count = 0

            while not done:
                actions = self.select_actions(observations)
                step_out = self.env.step(actions)
                next_observations = step_out["obs"]
                rewards = step_out["reward"]
                done = bool(step_out["terminated"] or step_out["truncated"])
                # Store the true-terminal flag (not truncation) so the Bellman
                # target keeps bootstrapping through timeout cut-offs; `done`
                # only ends the episode loop.
                terminal = bool(step_out["terminated"])

                for agent_id in self.agent_ids:
                    state = self._encode_observation(observations[agent_id])
                    next_state = self._encode_observation(next_observations[agent_id])
                    self._validate_state_shape(agent_id, state)
                    self._validate_state_shape(agent_id, next_state)

                    self.replay_buffers[agent_id].push(
                        state,
                        int(actions[agent_id]),
                        float(rewards[agent_id]),
                        next_state,
                        terminal,
                    )
                    episode_rewards[agent_id] += float(rewards[agent_id])
                    loss = self._optimize_agent(agent_id)
                    if loss is not None:
                        episode_losses[agent_id].append(loss)

                    if episode == 0 and step_count == 0 and self.debug_first_episode:
                        self._debug(
                            f"First transition | agent={agent_id} | "
                            f"state_shape={state.shape} | "
                            f"action={actions[agent_id]} | "
                            f"reward={float(rewards[agent_id]):.3f} | "
                            f"done={done}"
                        )

                observations = next_observations
                step_count += 1

            self.epsilon = max(self.min_epsilon, self.epsilon * self.epsilon_decay)

            if (episode + 1) % self.log_interval == 0:
                reward_str = ", ".join(
                    f"{aid}={value:.2f}" for aid, value in episode_rewards.items()
                )
                loss_str = ", ".join(
                    f"{aid}={np.mean(losses):.5f}" if losses else f"{aid}=warmup"
                    for aid, losses in episode_losses.items()
                )
                self._debug(
                    f"Episode {episode + 1}/{self.episodes} | eps={self.epsilon:.3f} | "
                    f"steps={step_count} | rewards: {reward_str} | avg_loss: {loss_str}"
                )

            if recorder:
                recorder.record(
                    episode + 1, self.epsilon, episode_rewards, episode_losses
                )

    # ------------------------------------------------------------------
    # Persistence
    # ------------------------------------------------------------------

    def save(self, path: str) -> None:
        os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
        payload = {
            "config": self.config,
            "agent_ids": self.agent_ids,
            "state_dim": self.state_dim,
            "action_dim": self.action_dim,
            "q_state_dicts": {
                aid: net.state_dict() for aid, net in self.q_networks.items()
            },
            "target_state_dicts": {
                aid: net.state_dict() for aid, net in self.target_networks.items()
            },
        }
        with open(path, "wb") as fh:
            pickle.dump(payload, fh)
        LOGGER.info("Saved DQN checkpoint -> %s", path)
        self._debug(f"Saved DQN checkpoint -> {path}")

    @classmethod
    def load(cls, env, config: dict, path: str) -> "DQN":
        instance = cls(env, config)
        with open(path, "rb") as fh:
            payload = pickle.load(fh)

        for agent_id in instance.agent_ids:
            if agent_id in payload["q_state_dicts"]:
                instance.q_networks[agent_id].load_state_dict(
                    payload["q_state_dicts"][agent_id]
                )
            if agent_id in payload["target_state_dicts"]:
                instance.target_networks[agent_id].load_state_dict(
                    payload["target_state_dicts"][agent_id]
                )

        LOGGER.info("Loaded DQN checkpoint from %s", path)
        return instance

ActorCritic

actor_critic.ActorCritic

Bases: BaseAlgorithm

Independent one-step online actor-critic: one network per agent, no replay buffer.

Source code in src/ppage/baselines/AC/actor_critic.py 
 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
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
class ActorCritic(BaseAlgorithm):
    """Independent one-step online actor-critic: one network per agent,
    no replay buffer."""

    def __init__(self, env, config: dict):
        super().__init__(env, config)

        self.gamma = float(config.get("gamma", 0.99))
        self.episodes = int(config.get("episodes", 1000))
        self.learning_rate = float(config.get("learning_rate", 1e-3))
        self.hidden_layers = [int(v) for v in config.get("hidden_layers", [128, 128])]
        self.value_coef = float(config.get("value_coef", 0.5))
        self.entropy_coef = float(config.get("entropy_coef", 0.0))
        self.grad_clip = float(config.get("grad_clip", 5.0))
        # L2 penalty on the policy head only (see A2C's actor_weight_decay,
        # docs/algorithms/a2c.md). AC's trunk is SHARED between the policy and
        # value heads, unlike A2C's fully separate actor/critic networks, so
        # decay is scoped to policy_head's own parameters via a dedicated
        # optimizer param group -- decaying the shared trunk would also
        # regularize the critic's features, not just the actor's. 0.0
        # preserves the original single-group behavior.
        self.actor_weight_decay = float(config.get("actor_weight_decay", 0.0))
        # Normalizes the critic's regression target (not just the reward) by
        # its own running mean/std -- fixes the actual root cause of the
        # entropy collapse actor_weight_decay only works around: with nothing
        # bounding the shared trunk's output scale, the critic's need to
        # represent this env's large-magnitude returns forces the trunk to
        # grow without limit, and policy_head's logits (a fixed linear
        # readout on those same trunk features) inherit that growth for
        # free, saturating the softmax independent of any actor gradient.
        # See docs/algorithms/actor-critic.md for the full diagnostic trail.
        # False preserves the original raw-value behavior (Sutton & Barto
        # Algorithm 13.5 as originally implemented).
        self.normalize_returns = bool(config.get("normalize_returns", False))
        self.return_norm_decay = float(config.get("return_norm_decay", 0.999))
        self.device = torch.device(config.get("device", "cpu"))
        self.verbose = bool(config.get("verbose", True))
        self.log_interval = int(config.get("log_interval", 10))
        self.debug_first_episode = bool(config.get("debug_first_episode", True))
        self.save_path = config.get("save_path", None)
        self.curves_path: Optional[str] = config.get("curves_path", None)

        seed = config.get("seed", None)
        self.rng = default_rng(seed)
        if seed is not None:
            torch.manual_seed(int(seed))

        # -- observation encoder contract (attached by build_environment) --
        self.observation_encoder = getattr(self.env, "observation_encoder", None)
        if self.observation_encoder is None:
            raise ValueError(
                "Environment is missing observation_encoder. Attach an observation "
                "builder's encode() method to env.observation_encoder before "
                "constructing ActorCritic (see run_from_config.build_environment)."
            )

        initial_obs, _ = self.env.reset()
        self.agent_ids = list(initial_obs.keys())
        self.state_dim = self._encode_observation(initial_obs[self.agent_ids[0]]).shape[
            0
        ]
        self.action_dim = self._resolve_action_dim(config)

        self._build_learners()

        self._debug(
            "Initialized ActorCritic | "
            f"agents={self.agent_ids} | state_dim={self.state_dim} | "
            f"action_dim={self.action_dim} | device={self.device} | "
            f"actor_weight_decay={self.actor_weight_decay} | "
            f"normalize_returns={self.normalize_returns}"
        )

    # ------------------------------------------------------------------
    # Setup helpers
    # ------------------------------------------------------------------

    def _resolve_action_dim(self, config: dict) -> int:
        """Infer action_dim from the env's action plugin, or validate if configured."""
        plugin = getattr(self.env, "action_space_plugin", None)
        plugin_n_actions = (
            int(plugin.n_actions)
            if plugin is not None
            else int(self.env.action_space.n)
        )

        configured = config.get("action_dim", None)
        if configured is None:
            return plugin_n_actions

        configured = int(configured)
        if configured != plugin_n_actions:
            raise ValueError(
                f"ActorCritic config error: 'action_dim'={configured} does not match "
                f"the environment's action space size ({plugin_n_actions}). Fix "
                "'action_dim' in your experiment YAML, or remove it entirely so "
                "it gets inferred automatically."
            )
        return configured

    def _build_learners(self) -> None:
        self.networks = {}
        self.optimizers = {}
        self.return_normalizers = {}

        for agent_id in self.agent_ids:
            network = ActorCriticNetwork(
                self.state_dim, self.hidden_layers, self.action_dim
            ).to(self.device)
            self.networks[agent_id] = network
            self.optimizers[agent_id] = optim.Adam(
                [
                    {
                        "params": network.policy_head.parameters(),
                        "weight_decay": self.actor_weight_decay,
                    },
                    {
                        "params": list(network.trunk.parameters())
                        + list(network.value_head.parameters()),
                    },
                ],
                lr=self.learning_rate,
            )
            if self.normalize_returns:
                self.return_normalizers[agent_id] = RunningReturnNormalizer(
                    decay=self.return_norm_decay
                )

    def _debug(self, message: str) -> None:
        if self.verbose:
            print(f"[ActorCritic] {message}")

    # ------------------------------------------------------------------
    # Observation encoding
    # ------------------------------------------------------------------

    def _encode_observation(self, obs) -> np.ndarray:
        encoded = self.observation_encoder(obs, self.env)
        return np.asarray(encoded, dtype=np.float32).reshape(-1)

    def _validate_state_shape(self, agent_id: str, state: np.ndarray) -> None:
        if state.shape[0] != self.state_dim:
            raise ValueError(
                f"ActorCritic expected state_dim={self.state_dim}, but agent "
                f"{agent_id!r} produced state_dim={state.shape[0]}. This usually "
                "means the observation plugin is returning variable-length "
                "observations."
            )

    # ------------------------------------------------------------------
    # Action selection (BaseAlgorithm interface)
    # ------------------------------------------------------------------

    def select_actions(self, observations: dict) -> dict:
        actions = {}
        for agent_id, obs in observations.items():
            state = self._encode_observation(obs)
            state_t = torch.from_numpy(state).float().unsqueeze(0).to(self.device)
            with torch.no_grad():
                logits, _ = self.networks[agent_id](state_t)
            dist = Categorical(logits=logits)
            actions[agent_id] = int(dist.sample().item())
        return actions

    # ------------------------------------------------------------------
    # Learning
    # ------------------------------------------------------------------

    def _update(
        self,
        agent_id: str,
        state: np.ndarray,
        action: int,
        reward: float,
        next_state: np.ndarray,
        terminal: bool,
        discount: float,
    ) -> tuple:
        """One TD(0) actor-critic step (Sutton & Barto, Algorithm 13.5).

        `terminal` reflects true termination only (not truncation), so the
        target keeps bootstrapping through timeout cut-offs -- same
        distinction DQN's training loop makes. `discount` is the running
        I <- gamma * I factor the caller decays once per step; it weights
        the actor's gradient by how many steps into the episode we are.
        """
        state_t = torch.from_numpy(state).float().unsqueeze(0).to(self.device)
        next_state_t = torch.from_numpy(next_state).float().unsqueeze(0).to(self.device)
        action_t = torch.tensor([action], device=self.device)

        with torch.no_grad():
            _, next_value = self.networks[agent_id](next_state_t)
            next_value = next_value * (0.0 if terminal else 1.0)

            if self.normalize_returns:
                # `value`/`next_value` are on a NORMALIZED scale once this is
                # enabled. Un-normalize next_value against the PRIOR estimate
                # to bootstrap in raw units, then update the running stats
                # AND analytically rescale value_head's weights (PopArt's
                # "preserving outputs precisely") so its denormalized
                # prediction stays consistent across this update -- see
                # return_normalizer.py's docstring for why the plain,
                # non-rescaling update() alone can diverge.
                normalizer = self.return_normalizers[agent_id]
                next_value_raw = normalizer.denormalize(next_value.item())
                raw_target = reward + self.gamma * next_value_raw
                normalizer.update_and_rescale_value_head(
                    raw_target, self.networks[agent_id].value_head
                )
                td_target = torch.tensor(
                    [normalizer.normalize(raw_target)],
                    dtype=torch.float32,
                    device=self.device,
                )
            else:
                td_target = reward + self.gamma * next_value

        # Forward pass for the CURRENT state happens AFTER any value_head
        # rescale above, so `value` and `td_target` are on the same
        # (freshly rescaled, if normalize_returns) normalized scale --
        # policy_head/trunk are untouched by the rescale, so this doesn't
        # change anything about action selection.
        logits, value = self.networks[agent_id](state_t)
        dist = Categorical(logits=logits)
        log_prob = dist.log_prob(action_t)
        entropy = dist.entropy()

        delta = td_target - value

        actor_loss = -(discount * delta.detach() * log_prob).mean()
        # Huber (SmoothL1), not raw squared error -- matches DQN's choice for
        # the same reason: robust to large TD errors instead of exploding
        # gradients when reward magnitudes are large. The actor's gradient
        # still uses the raw signed `delta` above; only the critic's own
        # regression loss changes.
        critic_loss = nn.SmoothL1Loss()(value, td_target)
        loss = (
            actor_loss
            + self.value_coef * critic_loss
            - self.entropy_coef * entropy.mean()
        )

        optimizer = self.optimizers[agent_id]
        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(
            self.networks[agent_id].parameters(), self.grad_clip
        )
        optimizer.step()

        return float(loss.item()), float(entropy.mean().item())

    # ------------------------------------------------------------------
    # Training loop (BaseAlgorithm interface)
    # ------------------------------------------------------------------

    def train(self):
        csv_file = None
        csv_writer = None
        if self.curves_path:
            os.makedirs(os.path.dirname(self.curves_path) or ".", exist_ok=True)
            csv_file = open(self.curves_path, "w", newline="")
            reward_cols = [f"{aid}_reward" for aid in self.agent_ids]
            loss_cols = [f"{aid}_loss" for aid in self.agent_ids]
            entropy_cols = [f"{aid}_entropy" for aid in self.agent_ids]
            csv_writer = csv.DictWriter(
                csv_file,
                fieldnames=["episode"] + reward_cols + loss_cols + entropy_cols,
            )
            csv_writer.writeheader()

        self._debug(f"Starting training for {self.episodes} episodes")
        try:
            self._train_loop(csv_writer)
        finally:
            if csv_file:
                csv_file.close()

        if self.save_path:
            self.save(self.save_path)

    def _train_loop(self, csv_writer) -> None:
        for episode in range(self.episodes):
            observations, _ = self.env.reset()
            discount = {agent_id: 1.0 for agent_id in self.agent_ids}
            episode_rewards = {agent_id: 0.0 for agent_id in self.agent_ids}
            episode_losses = {agent_id: [] for agent_id in self.agent_ids}
            episode_entropies = {agent_id: [] for agent_id in self.agent_ids}
            done = False
            step_count = 0

            while not done:
                actions = self.select_actions(observations)
                step_out = self.env.step(actions)
                next_observations = step_out["obs"]
                rewards = step_out["reward"]
                done = bool(step_out["terminated"] or step_out["truncated"])
                # True termination only -- keep bootstrapping through
                # timeout truncation, same distinction DQN's loop makes.
                terminal = bool(step_out["terminated"])

                for agent_id in self.agent_ids:
                    state = self._encode_observation(observations[agent_id])
                    next_state = self._encode_observation(next_observations[agent_id])
                    self._validate_state_shape(agent_id, state)
                    self._validate_state_shape(agent_id, next_state)

                    loss, entropy = self._update(
                        agent_id,
                        state,
                        int(actions[agent_id]),
                        float(rewards[agent_id]),
                        next_state,
                        terminal,
                        discount[agent_id],
                    )
                    discount[agent_id] *= self.gamma
                    episode_rewards[agent_id] += float(rewards[agent_id])
                    episode_losses[agent_id].append(loss)
                    episode_entropies[agent_id].append(entropy)

                    if episode == 0 and step_count == 0 and self.debug_first_episode:
                        self._debug(
                            f"First transition | agent={agent_id} | "
                            f"state_shape={state.shape} | "
                            f"action={actions[agent_id]} | "
                            f"reward={float(rewards[agent_id]):.3f} | "
                            f"done={done}"
                        )

                observations = next_observations
                step_count += 1

            if (episode + 1) % self.log_interval == 0:
                reward_str = ", ".join(
                    f"{aid}={value:.2f}" for aid, value in episode_rewards.items()
                )
                loss_str = ", ".join(
                    f"{aid}={np.mean(losses):.5f}"
                    for aid, losses in episode_losses.items()
                )
                entropy_str = ", ".join(
                    f"{aid}={np.mean(ent):.5f}"
                    for aid, ent in episode_entropies.items()
                )
                self._debug(
                    f"Episode {episode + 1}/{self.episodes} | "
                    f"steps={step_count} | rewards: {reward_str} | "
                    f"avg_loss: {loss_str} | avg_entropy: {entropy_str}"
                )

            if csv_writer:
                row: dict = {"episode": episode + 1}
                for aid in self.agent_ids:
                    row[f"{aid}_reward"] = round(episode_rewards[aid], 4)
                    losses = episode_losses[aid]
                    row[f"{aid}_loss"] = (
                        round(sum(losses) / len(losses), 6) if losses else ""
                    )
                    entropies = episode_entropies[aid]
                    row[f"{aid}_entropy"] = (
                        round(sum(entropies) / len(entropies), 6) if entropies else ""
                    )
                csv_writer.writerow(row)

    # ------------------------------------------------------------------
    # Persistence
    # ------------------------------------------------------------------

    def save(self, path: str) -> None:
        os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
        payload = {
            "config": self.config,
            "agent_ids": self.agent_ids,
            "state_dim": self.state_dim,
            "action_dim": self.action_dim,
            "state_dicts": {
                aid: net.state_dict() for aid, net in self.networks.items()
            },
        }
        if self.normalize_returns:
            # Persisted so a resumed/loaded run keeps interpreting the value
            # head's (normalized-scale) output consistently -- restarting the
            # normalizer from scratch would silently misinterpret an
            # already-trained value head's predictions.
            payload["return_normalizer_state"] = {
                aid: normalizer.state_dict()
                for aid, normalizer in self.return_normalizers.items()
            }
        with open(path, "wb") as fh:
            pickle.dump(payload, fh)
        LOGGER.info("Saved ActorCritic checkpoint -> %s", path)
        self._debug(f"Saved ActorCritic checkpoint -> {path}")

    @classmethod
    def load(cls, env, config: dict, path: str) -> "ActorCritic":
        instance = cls(env, config)
        with open(path, "rb") as fh:
            payload = pickle.load(fh)

        for agent_id in instance.agent_ids:
            if agent_id in payload["state_dicts"]:
                instance.networks[agent_id].load_state_dict(
                    payload["state_dicts"][agent_id]
                )

        if instance.normalize_returns and "return_normalizer_state" in payload:
            for agent_id in instance.agent_ids:
                if agent_id in payload["return_normalizer_state"]:
                    instance.return_normalizers[agent_id].load_state_dict(
                        payload["return_normalizer_state"][agent_id]
                    )

        LOGGER.info("Loaded ActorCritic checkpoint from %s", path)
        return instance

A2C

a2c.A2C

Bases: BaseAlgorithm

Independent A2C: one actor + one critic per agent, updated on an n-step rollout.

Source code in src/ppage/baselines/A2C/a2c.py 
 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
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
class A2C(BaseAlgorithm):
    """Independent A2C: one actor + one critic per agent, updated on an
    n-step rollout."""

    def __init__(self, env, config: dict):
        super().__init__(env, config)

        # ---- hyperparameters from config --------------------------------

        self.gamma = float(config.get("gamma", 0.99))
        self.episodes = int(config.get("episodes", 1000))
        self.hidden_layers = [int(v) for v in config.get("hidden_layers", [128, 128])]
        self.n_steps = int(config.get("n_steps", 5))
        self.entropy_coef = float(config.get("entropy_coef", 0.01))
        self.value_loss_coef = float(config.get("value_loss_coef", 0.5))
        self.grad_clip = float(config.get("grad_clip", 5.0))

        # separate learning rates: the critic usually wants to learn faster/
        # more stably than the actor, since the actor's gradient direction
        # depends on the critic already being roughly correct.
        actor_lr = config.get("actor_learning_rate", config.get("learning_rate", 3e-4))
        critic_lr = config.get(
            "critic_learning_rate", config.get("learning_rate", 1e-3)
        )
        self.actor_lr = float(actor_lr)
        self.critic_lr = float(critic_lr)
        # L2 penalty on the actor's weights -- discourages the logit
        # magnitude growth that drives the policy toward near-zero entropy
        # (confirmed empirically: max|logit| grows from ~3 to 40-60 over
        # training on this env, saturating the softmax; see
        # docs/algorithms/a2c.md). 0.0 preserves the original behavior.
        self.actor_weight_decay = float(config.get("actor_weight_decay", 0.0))

        self.device = torch.device(config.get("device", "cpu"))
        self.verbose = bool(config.get("verbose", True))
        self.log_interval = int(config.get("log_interval", 10))
        self.debug_first_episode = bool(config.get("debug_first_episode", True))
        self.save_path = config.get("save_path", None)
        self.curves_path: Optional[str] = config.get("curves_path", None)

        # If True, select_actions() acts greedily (argmax over the policy's
        # probabilities) instead of sampling. Exploration comes from sampling
        # the stochastic policy. Evaluation scripts set this to True.
        self.greedy_eval = bool(config.get("greedy_eval", False))

        if self.n_steps <= 0:
            raise ValueError(f"n_steps must be positive, got {self.n_steps}")

        seed = config.get("seed", None)
        self.rng = default_rng(seed)
        if seed is not None:
            torch.manual_seed(int(seed))

        # ---- observation encoder contract --------

        self.observation_encoder = getattr(self.env, "observation_encoder", None)
        if self.observation_encoder is None:
            raise ValueError(
                "Environment is missing observation_encoder. Attach an observation "
                "builder's encode() method to env.observation_encoder before "
                "constructing A2C (see run_from_config.build_environment)."
            )

        initial_obs, _ = self.env.reset()
        self.agent_ids = list(initial_obs.keys())
        self.state_dim = self._encode_observation(initial_obs[self.agent_ids[0]]).shape[
            0
        ]
        self.action_dim = self._resolve_action_dim(config)

        self._build_learners()
        self._train_steps = 0

        self._debug(
            "Initialized A2C | "
            f"agents={self.agent_ids} | state_dim={self.state_dim} | "
            f"action_dim={self.action_dim} | device={self.device} | "
            f"n_steps={self.n_steps} | entropy_coef={self.entropy_coef} | "
            f"actor_weight_decay={self.actor_weight_decay}"
        )

    # ------------------------------------------------------------------
    # Setup helpers
    # ------------------------------------------------------------------

    def _resolve_action_dim(self, config: dict) -> int:
        """Infer action_dim from the env's action plugin, or validate it if configured.
        Identical logic to DQN's, kept local so A2C stays self-contained."""
        plugin = getattr(self.env, "action_space_plugin", None)
        plugin_n_actions = (
            int(plugin.n_actions)
            if plugin is not None
            else int(self.env.action_space.n)
        )

        configured = config.get("action_dim", None)
        if configured is None:
            return plugin_n_actions

        configured = int(configured)
        if configured != plugin_n_actions:
            raise ValueError(
                f"A2C config error: 'action_dim'={configured} does not match the "
                f"environment's action space size ({plugin_n_actions}). Fix "
                "'action_dim' in your experiment YAML, or remove it entirely so "
                "it gets inferred automatically."
            )
        return configured

    def _build_learners(self) -> None:
        self.actors = {}
        self.critics = {}
        self.actor_optimizers = {}
        self.critic_optimizers = {}

        for agent_id in self.agent_ids:
            self.actors[agent_id] = ActorNetwork(
                self.state_dim, self.hidden_layers, self.action_dim
            ).to(self.device)
            self.critics[agent_id] = CriticNetwork(
                self.state_dim, self.hidden_layers
            ).to(self.device)

            self.actor_optimizers[agent_id] = optim.Adam(
                self.actors[agent_id].parameters(),
                lr=self.actor_lr,
                weight_decay=self.actor_weight_decay,
            )
            self.critic_optimizers[agent_id] = optim.Adam(
                self.critics[agent_id].parameters(), lr=self.critic_lr
            )

    def _debug(self, message: str) -> None:
        if self.verbose:
            print(f"[A2C] {message}")

    def _empty_rollout(self) -> dict:
        return {
            aid: {"log_probs": [], "values": [], "rewards": [], "entropies": []}
            for aid in self.agent_ids
        }

    # ------------------------------------------------------------------
    # Observation encoding
    # ------------------------------------------------------------------

    def _encode_observation(self, obs) -> np.ndarray:
        encoded = self.observation_encoder(obs, self.env)
        return np.asarray(encoded, dtype=np.float32).reshape(-1)

    def _validate_state_shape(self, agent_id: str, state: np.ndarray) -> None:
        if state.shape[0] != self.state_dim:
            raise ValueError(
                f"A2C expected state_dim={self.state_dim}, but agent {agent_id!r} "
                f"produced state_dim={state.shape[0]}. This usually means the "
                "observation plugin is returning variable-length observations."
            )

    # ------------------------------------------------------------------
    # Action selection (BaseAlgorithm interface)
    # ------------------------------------------------------------------

    def select_actions(self, observations: dict) -> dict:
        """
        Public, gradient-free action selection -- used by BaseAlgorithm.evaluate()
        and by anything outside the training loop. The training loop itself does
        NOT call this method: it needs log_prob/value/entropy WITH gradients
        attached, so it runs the same forward pass directly (see _train_loop).
        """
        actions = {}
        with torch.no_grad():
            for agent_id, obs in observations.items():
                state = self._encode_observation(obs)
                state_t = torch.from_numpy(state).float().unsqueeze(0).to(self.device)
                dist = self.actors[agent_id].distribution(state_t)
                if self.greedy_eval:
                    action = torch.argmax(dist.probs, dim=1)
                else:
                    action = dist.sample()
                actions[agent_id] = int(action.item())
        return actions

    # ------------------------------------------------------------------
    # Learning
    # ------------------------------------------------------------------

    def _optimize_agent(self, agent_id: str, traj: dict, bootstrap_value: float):
        """
        Turns one agent's n-step rollout into an actor update and a critic
        update. Returns (actor_loss, critic_loss, mean_entropy) as floats for
        logging.

        n-step return:  R_t = r_t + gamma*r_{t+1} + ... + gamma^k * V(s_{t+k})
        Bootstrap value is 0 if the episode ended inside this rollout
        (nothing beyond a terminal state contributes future reward),
        otherwise it's the critic's own estimate of the state we stopped at.
        """
        rewards = traj["rewards"]

        returns = []
        R = bootstrap_value
        for r in reversed(rewards):
            R = r + self.gamma * R
            returns.insert(0, R)
        returns_t = torch.tensor(returns, dtype=torch.float32, device=self.device)

        values_t = torch.cat(traj["values"])  # keeps grad -> critic params
        log_probs_t = torch.cat(traj["log_probs"])  # keeps grad -> actor params
        entropies_t = torch.cat(traj["entropies"])

        # detach() is essential:
        # the actor should be pushed by the SIGN and SIZE of the advantage,
        # not by trying to change the critic's value through this path.
        advantages = returns_t - values_t.detach()

        actor_loss = (
            -(log_probs_t * advantages).mean() - self.entropy_coef * entropies_t.mean()
        )
        # Huber (SmoothL1), not raw MSE -- matches AC/DQN's choice: robust to
        # large TD errors instead of exploding gradients when this env's
        # accumulated reward magnitudes run into the thousands per episode.
        critic_loss = nn.SmoothL1Loss()(values_t, returns_t)

        self.actor_optimizers[agent_id].zero_grad()
        actor_loss.backward()
        torch.nn.utils.clip_grad_norm_(
            self.actors[agent_id].parameters(), self.grad_clip
        )
        self.actor_optimizers[agent_id].step()

        self.critic_optimizers[agent_id].zero_grad()
        (self.value_loss_coef * critic_loss).backward()
        torch.nn.utils.clip_grad_norm_(
            self.critics[agent_id].parameters(), self.grad_clip
        )
        self.critic_optimizers[agent_id].step()

        self._train_steps += 1

        return (
            float(actor_loss.item()),
            float(critic_loss.item()),
            float(entropies_t.mean().item()),
        )

    # ------------------------------------------------------------------
    # Training loop (BaseAlgorithm interface)
    # ------------------------------------------------------------------

    def train(self):
        csv_file = None
        csv_writer = None
        if self.curves_path:
            os.makedirs(os.path.dirname(self.curves_path) or ".", exist_ok=True)
            csv_file = open(self.curves_path, "w", newline="")
            reward_cols = [f"{aid}_reward" for aid in self.agent_ids]
            actor_loss_cols = [f"{aid}_actor_loss" for aid in self.agent_ids]
            critic_loss_cols = [f"{aid}_critic_loss" for aid in self.agent_ids]
            entropy_cols = [f"{aid}_entropy" for aid in self.agent_ids]
            csv_writer = csv.DictWriter(
                csv_file,
                fieldnames=["episode"]
                + reward_cols
                + actor_loss_cols
                + critic_loss_cols
                + entropy_cols,
            )
            csv_writer.writeheader()

        self._debug(f"Starting training for {self.episodes} episodes")
        try:
            self._train_loop(csv_writer)
        finally:
            if csv_file:
                csv_file.close()

        if self.save_path:
            self.save(self.save_path)

    def _act_with_grad(self, agent_id: str, obs) -> tuple:
        """Forward pass WITH gradient tracking -- used only inside the
        training loop, where log_prob/value/entropy must stay part of the
        autograd graph so _optimize_agent can backpropagate through them."""
        state = self._encode_observation(obs)
        self._validate_state_shape(agent_id, state)
        state_t = torch.from_numpy(state).float().unsqueeze(0).to(self.device)

        dist = self.actors[agent_id].distribution(state_t)
        action = dist.sample()
        log_prob = dist.log_prob(action)
        entropy = dist.entropy()
        value = self.critics[agent_id](state_t).squeeze(1)

        return int(action.item()), log_prob, entropy, value

    def _train_loop(self, csv_writer) -> None:
        for episode in range(self.episodes):
            observations, _ = self.env.reset()
            episode_rewards = {aid: 0.0 for aid in self.agent_ids}
            episode_actor_losses = {aid: [] for aid in self.agent_ids}
            episode_critic_losses = {aid: [] for aid in self.agent_ids}
            episode_entropies = {aid: [] for aid in self.agent_ids}

            rollout = self._empty_rollout()
            done = False
            step_count = 0

            while not done:
                actions = {}
                for agent_id in self.agent_ids:
                    action, log_prob, entropy, value = self._act_with_grad(
                        agent_id, observations[agent_id]
                    )
                    actions[agent_id] = action
                    rollout[agent_id]["log_probs"].append(log_prob)
                    rollout[agent_id]["values"].append(value)
                    rollout[agent_id]["entropies"].append(entropy)

                step_out = self.env.step(actions)
                next_observations = step_out["obs"]
                rewards = step_out["reward"]
                done = bool(step_out["terminated"] or step_out["truncated"])

                for agent_id in self.agent_ids:
                    r = float(rewards[agent_id])
                    rollout[agent_id]["rewards"].append(r)
                    episode_rewards[agent_id] += r

                    if episode == 0 and step_count == 0 and self.debug_first_episode:
                        self._debug(
                            f"First transition | agent={agent_id} | "
                            f"action={actions[agent_id]} | reward={r:.3f} | done={done}"
                        )

                observations = next_observations
                step_count += 1

                # Update every n_steps, or immediately at episode end (bootstraps
                # from 0 instead of from V(s_next)).
                ready_to_update = (step_count % self.n_steps == 0) or done
                if ready_to_update:
                    for agent_id in self.agent_ids:
                        if not rollout[agent_id]["rewards"]:
                            continue  # nothing collected for this agent yet

                        if done:
                            bootstrap_value = 0.0
                        else:
                            with torch.no_grad():
                                next_state = self._encode_observation(
                                    observations[agent_id]
                                )
                                next_state_t = (
                                    torch.from_numpy(next_state)
                                    .float()
                                    .unsqueeze(0)
                                    .to(self.device)
                                )
                                bootstrap_value = float(
                                    self.critics[agent_id](next_state_t).item()
                                )

                        actor_loss, critic_loss, mean_entropy = self._optimize_agent(
                            agent_id, rollout[agent_id], bootstrap_value
                        )
                        episode_actor_losses[agent_id].append(actor_loss)
                        episode_critic_losses[agent_id].append(critic_loss)
                        episode_entropies[agent_id].append(mean_entropy)

                    rollout = self._empty_rollout()

            if (episode + 1) % self.log_interval == 0:
                reward_str = ", ".join(
                    f"{aid}={value:.2f}" for aid, value in episode_rewards.items()
                )
                critic_loss_str = ", ".join(
                    f"{aid}={np.mean(losses):.5f}" if losses else f"{aid}=warmup"
                    for aid, losses in episode_critic_losses.items()
                )
                actor_loss_str = ", ".join(
                    f"{aid}={np.mean(losses):.5f}" if losses else f"{aid}=warmup"
                    for aid, losses in episode_actor_losses.items()
                )
                entropy_str = ", ".join(
                    f"{aid}={np.mean(ent):.5f}" if ent else f"{aid}=warmup"
                    for aid, ent in episode_entropies.items()
                )
                self._debug(
                    f"Episode {episode + 1}/{self.episodes} | steps={step_count} | "
                    f"rewards: {reward_str} | avg_critic_loss: {critic_loss_str} | "
                    f"avg_actor_loss: {actor_loss_str} | avg_entropy: {entropy_str}"
                )

            if csv_writer:
                row: dict = {"episode": episode + 1}
                for aid in self.agent_ids:
                    row[f"{aid}_reward"] = round(episode_rewards[aid], 4)
                    a_losses = episode_actor_losses[aid]
                    c_losses = episode_critic_losses[aid]
                    ent = episode_entropies[aid]
                    row[f"{aid}_actor_loss"] = (
                        round(float(np.mean(a_losses)), 6) if a_losses else ""
                    )
                    row[f"{aid}_critic_loss"] = (
                        round(float(np.mean(c_losses)), 6) if c_losses else ""
                    )
                    row[f"{aid}_entropy"] = round(float(np.mean(ent)), 6) if ent else ""
                csv_writer.writerow(row)

    # ------------------------------------------------------------------
    # Persistence
    # ------------------------------------------------------------------

    def save(self, path: str) -> None:
        os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
        payload = {
            "config": self.config,
            "agent_ids": self.agent_ids,
            "state_dim": self.state_dim,
            "action_dim": self.action_dim,
            "actor_state_dicts": {
                aid: net.state_dict() for aid, net in self.actors.items()
            },
            "critic_state_dicts": {
                aid: net.state_dict() for aid, net in self.critics.items()
            },
        }
        with open(path, "wb") as fh:
            pickle.dump(payload, fh)
        LOGGER.info("Saved A2C checkpoint -> %s", path)
        self._debug(f"Saved A2C checkpoint -> {path}")

    @classmethod
    def load(cls, env, config: dict, path: str) -> "A2C":
        instance = cls(env, config)
        with open(path, "rb") as fh:
            payload = pickle.load(fh)

        for agent_id in instance.agent_ids:
            if agent_id in payload["actor_state_dicts"]:
                instance.actors[agent_id].load_state_dict(
                    payload["actor_state_dicts"][agent_id]
                )
            if agent_id in payload["critic_state_dicts"]:
                instance.critics[agent_id].load_state_dict(
                    payload["critic_state_dicts"][agent_id]
                )

        LOGGER.info("Loaded A2C checkpoint from %s", path)
        return instance

select_actions 

select_actions(observations: dict) -> dict

Public, gradient-free action selection -- used by BaseAlgorithm.evaluate() and by anything outside the training loop. The training loop itself does NOT call this method: it needs log_prob/value/entropy WITH gradients attached, so it runs the same forward pass directly (see _train_loop).

Source code in src/ppage/baselines/A2C/a2c.py 
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
def select_actions(self, observations: dict) -> dict:
    """
    Public, gradient-free action selection -- used by BaseAlgorithm.evaluate()
    and by anything outside the training loop. The training loop itself does
    NOT call this method: it needs log_prob/value/entropy WITH gradients
    attached, so it runs the same forward pass directly (see _train_loop).
    """
    actions = {}
    with torch.no_grad():
        for agent_id, obs in observations.items():
            state = self._encode_observation(obs)
            state_t = torch.from_numpy(state).float().unsqueeze(0).to(self.device)
            dist = self.actors[agent_id].distribution(state_t)
            if self.greedy_eval:
                action = torch.argmax(dist.probs, dim=1)
            else:
                action = dist.sample()
            actions[agent_id] = int(action.item())
    return actions

A3C

a3c.A3C

Bases: BaseAlgorithm

Independent A3C: one shared network + shared optimizer per agent, updated asynchronously by multiple worker processes.

Source code in src/ppage/baselines/A3C/a3c.py 
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
class A3C(BaseAlgorithm):
    """Independent A3C: one shared network + shared optimizer per agent,
    updated asynchronously by multiple worker processes."""

    def __init__(self, env, config: dict):
        super().__init__(env, config)

        self.gamma = float(config.get("gamma", 0.99))
        self.episodes = int(config.get("episodes", 1000))
        self.learning_rate = float(config.get("learning_rate", 1e-3))
        self.hidden_layers = [int(v) for v in config.get("hidden_layers", [128, 128])]
        self.n_steps = int(config.get("n_steps", 5))
        self.entropy_coef = float(config.get("entropy_coef", 0.05))
        self.value_coef = float(config.get("value_coef", 0.5))
        self.grad_clip = float(config.get("grad_clip", 5.0))
        # L2 penalty on the policy head only (see A2C's actor_weight_decay,
        # docs/algorithms/a2c.md). A3C reuses AC's shared-trunk network, so
        # decay is scoped to policy_head's own parameters via a dedicated
        # optimizer param group -- decaying the shared trunk would also
        # regularize the critic's features across every worker, not just the
        # actor's. 0.0 preserves the original single-group behavior.
        self.actor_weight_decay = float(config.get("actor_weight_decay", 0.0))
        # See ppage/baselines/AC/return_normalizer.py and AC's own normalize_returns
        # (docs/algorithms/actor-critic.md) -- same fix, using ONE shared,
        # lock-protected SharedReturnNormalizer per agent across all workers
        # (see _build_learners/_worker_loop), so the PopArt weight-rescale
        # applies consistently instead of workers fighting each other.
        self.normalize_returns = bool(config.get("normalize_returns", False))
        self.return_norm_decay = float(config.get("return_norm_decay", 0.999))
        self.num_workers = int(config.get("num_workers", 4))
        self.verbose = bool(config.get("verbose", True))
        self.log_interval = int(config.get("log_interval", 10))
        self.save_path = config.get("save_path", None)
        self.curves_path: Optional[str] = config.get("curves_path", None)

        if self.n_steps <= 0:
            raise ValueError(f"n_steps must be positive, got {self.n_steps}")
        if self.num_workers <= 0:
            raise ValueError(f"num_workers must be positive, got {self.num_workers}")

        self.env_fn = config.get("env_fn", None)
        if self.env_fn is None or not callable(self.env_fn):
            raise ValueError(
                "A3C requires config['env_fn']: a zero-argument, picklable "
                "callable that builds a fresh environment for each worker "
                "process (a lambda will NOT pickle under the 'spawn' start "
                "method -- use a module-level function or a callable class "
                "instance; see scripts/run_a3c.py's EnvFactory)."
            )

        self.seed = config.get("seed", None)
        if self.seed is not None:
            torch.manual_seed(int(self.seed))

        self.observation_encoder = getattr(self.env, "observation_encoder", None)
        if self.observation_encoder is None:
            raise ValueError(
                "Environment is missing observation_encoder. Attach an observation "
                "builder's encode() method to env.observation_encoder before "
                "constructing A3C (see run_from_config.build_environment)."
            )

        initial_obs, _ = self.env.reset()
        self.agent_ids = list(initial_obs.keys())
        self.state_dim = self._encode_observation(initial_obs[self.agent_ids[0]]).shape[
            0
        ]
        self.action_dim = self._resolve_action_dim(config)

        self._build_learners()

        self._debug(
            "Initialized A3C | "
            f"agents={self.agent_ids} | state_dim={self.state_dim} | "
            f"action_dim={self.action_dim} | num_workers={self.num_workers} | "
            f"n_steps={self.n_steps} | entropy_coef={self.entropy_coef} | "
            f"actor_weight_decay={self.actor_weight_decay} | "
            f"normalize_returns={self.normalize_returns} | device=cpu"
        )

    # ------------------------------------------------------------------
    # Setup helpers
    # ------------------------------------------------------------------

    def _resolve_action_dim(self, config: dict) -> int:
        plugin = getattr(self.env, "action_space_plugin", None)
        plugin_n_actions = (
            int(plugin.n_actions)
            if plugin is not None
            else int(self.env.action_space.n)
        )

        configured = config.get("action_dim", None)
        if configured is None:
            return plugin_n_actions

        configured = int(configured)
        if configured != plugin_n_actions:
            raise ValueError(
                f"A3C config error: 'action_dim'={configured} does not match the "
                f"environment's action space size ({plugin_n_actions}). Fix "
                "'action_dim' in your experiment YAML, or remove it entirely so "
                "it gets inferred automatically."
            )
        return configured

    def _build_learners(self) -> None:
        self.networks = {}
        self.optimizers = {}
        self.return_normalizers = {}

        for agent_id in self.agent_ids:
            network = ActorCriticNetwork(
                self.state_dim, self.hidden_layers, self.action_dim
            )
            network.share_memory()
            self.networks[agent_id] = network
            self.optimizers[agent_id] = SharedAdam(
                [
                    {
                        "params": network.policy_head.parameters(),
                        "weight_decay": self.actor_weight_decay,
                    },
                    {
                        "params": list(network.trunk.parameters())
                        + list(network.value_head.parameters()),
                    },
                ],
                lr=self.learning_rate,
            )
            if self.normalize_returns:
                # ONE shared, lock-protected estimate per agent, created
                # here in the main process and passed to every worker --
                # not one-per-worker-local -- so the PopArt weight-rescale
                # below has a single consistent (mean, std) to stay
                # consistent with. See return_normalizer.py's
                # SharedReturnNormalizer docstring for why.
                self.return_normalizers[agent_id] = SharedReturnNormalizer(
                    decay=self.return_norm_decay
                )

    def _debug(self, message: str) -> None:
        if self.verbose:
            print(f"[A3C] {message}")

    def _encode_observation(self, obs) -> np.ndarray:
        encoded = self.observation_encoder(obs, self.env)
        return np.asarray(encoded, dtype=np.float32).reshape(-1)

    # ------------------------------------------------------------------
    # Action selection (BaseAlgorithm interface) -- main process only,
    # against the shared global networks (used by evaluate()).
    # ------------------------------------------------------------------

    def select_actions(self, observations: dict) -> dict:
        actions = {}
        for agent_id, obs in observations.items():
            state = self._encode_observation(obs)
            state_t = torch.from_numpy(state).float().unsqueeze(0)
            with torch.no_grad():
                logits, _ = self.networks[agent_id](state_t)
            dist = Categorical(logits=logits)
            actions[agent_id] = int(dist.sample().item())
        return actions

    # ------------------------------------------------------------------
    # Training (BaseAlgorithm interface)
    # ------------------------------------------------------------------

    def train(self):
        if self.curves_path:
            os.makedirs(os.path.dirname(self.curves_path) or ".", exist_ok=True)
            reward_cols = [f"{aid}_reward" for aid in self.agent_ids]
            loss_cols = [f"{aid}_loss" for aid in self.agent_ids]
            entropy_cols = [f"{aid}_entropy" for aid in self.agent_ids]
            with open(self.curves_path, "w", newline="") as f:
                writer = csv.DictWriter(
                    f,
                    fieldnames=["episode", "worker"]
                    + reward_cols
                    + loss_cols
                    + entropy_cols,
                )
                writer.writeheader()

        self._debug(
            f"Starting training for {self.episodes} episodes across "
            f"{self.num_workers} workers"
        )

        episode_counter = mp.Value("i", 0)
        counter_lock = mp.Lock()
        csv_lock = mp.Lock() if self.curves_path else None

        processes = []
        for worker_id in range(self.num_workers):
            worker_seed = None if self.seed is None else int(self.seed) + worker_id
            p = mp.Process(
                target=_worker_loop,
                args=(
                    worker_id,
                    self.agent_ids,
                    self.state_dim,
                    self.action_dim,
                    self.hidden_layers,
                    self.networks,
                    self.optimizers,
                    self.env_fn,
                    self.gamma,
                    self.n_steps,
                    self.entropy_coef,
                    self.value_coef,
                    self.grad_clip,
                    episode_counter,
                    self.episodes,
                    counter_lock,
                    csv_lock,
                    self.curves_path,
                    self.log_interval,
                    self.verbose,
                    worker_seed,
                    self.return_normalizers if self.normalize_returns else None,
                ),
            )
            p.start()
            processes.append(p)

        for p in processes:
            p.join()

        if self.save_path:
            self.save(self.save_path)

    # ------------------------------------------------------------------
    # Persistence
    # ------------------------------------------------------------------

    def save(self, path: str) -> None:
        os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
        # env_fn is a live callable tied to this session, not meaningful
        # checkpoint data -- and may not even be picklable (e.g. a lambda),
        # so it's dropped before persisting rather than saved wholesale.
        config_to_save = {k: v for k, v in self.config.items() if k != "env_fn"}
        payload = {
            "config": config_to_save,
            "agent_ids": self.agent_ids,
            "state_dim": self.state_dim,
            "action_dim": self.action_dim,
            "state_dicts": {
                aid: net.state_dict() for aid, net in self.networks.items()
            },
        }
        if self.normalize_returns:
            # Persisted so a resumed/loaded run keeps interpreting the
            # shared value_head's (normalized-scale) output consistently --
            # restarting the normalizer from scratch would silently
            # misinterpret an already-trained value_head's predictions.
            # Matches ActorCritic.save()'s equivalent.
            payload["return_normalizer_state"] = {
                aid: normalizer.state_dict()
                for aid, normalizer in self.return_normalizers.items()
            }
        with open(path, "wb") as fh:
            pickle.dump(payload, fh)
        LOGGER.info("Saved A3C checkpoint -> %s", path)
        self._debug(f"Saved A3C checkpoint -> {path}")

    @classmethod
    def load(cls, env, config: dict, path: str) -> "A3C":
        instance = cls(env, config)
        with open(path, "rb") as fh:
            payload = pickle.load(fh)

        for agent_id in instance.agent_ids:
            if agent_id in payload["state_dicts"]:
                instance.networks[agent_id].load_state_dict(
                    payload["state_dicts"][agent_id]
                )

        if instance.normalize_returns and "return_normalizer_state" in payload:
            for agent_id in instance.agent_ids:
                if agent_id in payload["return_normalizer_state"]:
                    instance.return_normalizers[agent_id].load_state_dict(
                        payload["return_normalizer_state"][agent_id]
                    )

        LOGGER.info("Loaded A3C checkpoint from %s", path)
        return instance

SharedAdam

shared_adam.SharedAdam

Bases: Adam

Source code in src/ppage/baselines/A3C/shared_adam.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
class SharedAdam(torch.optim.Adam):
    def __init__(
        self,
        params,
        lr: float = 1e-3,
        betas: tuple = (0.9, 0.999),
        eps: float = 1e-8,
        weight_decay: float = 0,
    ):
        super().__init__(params, lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
        for group in self.param_groups:
            for p in group["params"]:
                state = self.state[p]
                state["step"] = torch.zeros(1)
                state["exp_avg"] = torch.zeros_like(p.data)
                state["exp_avg_sq"] = torch.zeros_like(p.data)
                state["step"].share_memory_()
                state["exp_avg"].share_memory_()
                state["exp_avg_sq"].share_memory_()