CLRF

.gitattributes

@@ -0,0 +1,13 @@
+# Shell scripts must always use LF (Linux) line endings
+*.sh text eol=lf
+
+# Windows-specific scripts
+*.ps1 text eol=crlf
+*.bat text eol=crlf
+*.cmd text eol=crlf
+
+# YAML and Docker files should be LF
+*.yml text eol=lf
+*.yaml text eol=lf
+Dockerfile text eol=lf
+.dockerignore text eol=lf

deploy/grafana/provisioning/dashboards/json/antiatropos-overview.json

@@ -1,647 +1,647 @@
-{
-  "annotations": {
-    "list": [
-      {
-        "builtIn": 1,
-        "datasource": {
-          "type": "grafana",
-          "uid": "-- Grafana --"
-        },
-        "enable": true,
-        "hide": true,
-        "iconColor": "rgba(0, 211, 255, 1)",
-        "name": "Annotations & Alerts",
-        "type": "dashboard"
-      }
-    ]
-  },
-  "editable": true,
-  "fiscalYearStartMonth": 0,
-  "graphTooltip": 1,
-  "id": null,
-  "links": [],
-  "liveNow": false,
-  "panels": [
-    {
-      "datasource": {
-        "type": "prometheus",
-        "uid": "PBFA97CFB590B2093"
-      },
-      "fieldConfig": {
-        "defaults": {
-          "color": {
-            "mode": "thresholds"
-          },
-          "decimals": 3,
-          "mappings": [],
-          "thresholds": {
-            "mode": "absolute",
-            "steps": [
-              {
-                "color": "green",
-                "value": null
-              },
-              {
-                "color": "red",
-                "value": 0
-              }
-            ]
-          },
-          "unit": "short"
-        },
-        "overrides": []
-      },
-      "gridPos": {
-        "h": 4,
-        "w": 6,
-        "x": 0,
-        "y": 0
-      },
-      "id": 1,
-      "options": {
-        "colorMode": "value",
-        "graphMode": "none",
-        "justifyMode": "auto",
-        "orientation": "auto",
-        "reduceOptions": {
-          "calcs": [
-            "lastNotNull"
-          ],
-          "fields": "",
-          "values": false
-        },
-        "textMode": "auto"
-      },
-      "pluginVersion": "12.3.1",
-      "targets": [
-        {
-          "editorMode": "code",
+{
+  "annotations": {
+    "list": [
+      {
+        "builtIn": 1,
+        "datasource": {
+          "type": "grafana",
+          "uid": "-- Grafana --"
+        },
+        "enable": true,
+        "hide": true,
+        "iconColor": "rgba(0, 211, 255, 1)",
+        "name": "Annotations & Alerts",
+        "type": "dashboard"
+      }
+    ]
+  },
+  "editable": true,
+  "fiscalYearStartMonth": 0,
+  "graphTooltip": 1,
+  "id": null,
+  "links": [],
+  "liveNow": false,
+  "panels": [
+    {
+      "datasource": {
+        "type": "prometheus",
+        "uid": "PBFA97CFB590B2093"
+      },
+      "fieldConfig": {
+        "defaults": {
+          "color": {
+            "mode": "thresholds"
+          },
+          "decimals": 3,
+          "mappings": [],
+          "thresholds": {
+            "mode": "absolute",
+            "steps": [
+              {
+                "color": "green",
+                "value": null
+              },
+              {
+                "color": "red",
+                "value": 0
+              }
+            ]
+          },
+          "unit": "short"
+        },
+        "overrides": []
+      },
+      "gridPos": {
+        "h": 4,
+        "w": 6,
+        "x": 0,
+        "y": 0
+      },
+      "id": 1,
+      "options": {
+        "colorMode": "value",
+        "graphMode": "none",
+        "justifyMode": "auto",
+        "orientation": "auto",
+        "reduceOptions": {
+          "calcs": [
+            "lastNotNull"
+          ],
+          "fields": "",
+          "values": false
+        },
+        "textMode": "auto"
+      },
+      "pluginVersion": "12.3.1",
+      "targets": [
+        {
+          "editorMode": "code",
           "expr": "scalar(avg(last_over_time(antiatropos_reward{mode=~\"live|simulated|hybrid|aws\"}[1m])))",
           "legendFormat": "reward (all modes)",
-          "range": true,
-          "refId": "A"
-        }
-      ],
-      "title": "Latest Reward",
-      "type": "stat"
-    },
-    {
-      "datasource": {
-        "type": "prometheus",
-        "uid": "PBFA97CFB590B2093"
-      },
-      "fieldConfig": {
-        "defaults": {
-          "color": {
-            "mode": "thresholds"
-          },
-          "decimals": 3,
-          "mappings": [],
-          "thresholds": {
-            "mode": "absolute",
-            "steps": [
-              {
-                "color": "green",
-                "value": null
-              },
-              {
-                "color": "orange",
-                "value": 0.4
-              },
-              {
-                "color": "red",
-                "value": 0.8
-              }
-            ]
-          },
-          "unit": "short"
-        },
-        "overrides": []
-      },
-      "gridPos": {
-        "h": 4,
-        "w": 6,
-        "x": 6,
-        "y": 0
-      },
-      "id": 2,
-      "options": {
-        "colorMode": "value",
-        "graphMode": "none",
-        "justifyMode": "auto",
-        "orientation": "auto",
-        "reduceOptions": {
-          "calcs": [
-            "lastNotNull"
-          ],
-          "fields": "",
-          "values": false
-        },
-        "textMode": "auto"
-      },
-      "pluginVersion": "12.3.1",
-      "targets": [
-        {
-          "editorMode": "code",
+          "range": true,
+          "refId": "A"
+        }
+      ],
+      "title": "Latest Reward",
+      "type": "stat"
+    },
+    {
+      "datasource": {
+        "type": "prometheus",
+        "uid": "PBFA97CFB590B2093"
+      },
+      "fieldConfig": {
+        "defaults": {
+          "color": {
+            "mode": "thresholds"
+          },
+          "decimals": 3,
+          "mappings": [],
+          "thresholds": {
+            "mode": "absolute",
+            "steps": [
+              {
+                "color": "green",
+                "value": null
+              },
+              {
+                "color": "orange",
+                "value": 0.4
+              },
+              {
+                "color": "red",
+                "value": 0.8
+              }
+            ]
+          },
+          "unit": "short"
+        },
+        "overrides": []
+      },
+      "gridPos": {
+        "h": 4,
+        "w": 6,
+        "x": 6,
+        "y": 0
+      },
+      "id": 2,
+      "options": {
+        "colorMode": "value",
+        "graphMode": "none",
+        "justifyMode": "auto",
+        "orientation": "auto",
+        "reduceOptions": {
+          "calcs": [
+            "lastNotNull"
+          ],
+          "fields": "",
+          "values": false
+        },
+        "textMode": "auto"
+      },
+      "pluginVersion": "12.3.1",
+      "targets": [
+        {
+          "editorMode": "code",
           "expr": "scalar(avg(last_over_time(antiatropos_total_queue_backlog{mode=~\"live|simulated|hybrid|aws\"}[1m])))",
           "legendFormat": "queue backlog (all modes)",
-          "range": true,
-          "refId": "A"
-        }
-      ],
-      "title": "Queue Backlog (Norm)",
-      "type": "stat"
-    },
-    {
-      "datasource": {
-        "type": "prometheus",
-        "uid": "PBFA97CFB590B2093"
-      },
-      "fieldConfig": {
-        "defaults": {
-          "color": {
-            "mode": "thresholds"
-          },
-          "decimals": 3,
-          "mappings": [],
-          "thresholds": {
-            "mode": "absolute",
-            "steps": [
-              {
-                "color": "green",
-                "value": null
-              },
-              {
-                "color": "orange",
-                "value": 0.4
-              },
-              {
-                "color": "red",
-                "value": 0.8
-              }
-            ]
-          },
-          "unit": "short"
-        },
-        "overrides": []
-      },
-      "gridPos": {
-        "h": 4,
-        "w": 6,
-        "x": 12,
-        "y": 0
-      },
-      "id": 3,
-      "options": {
-        "colorMode": "value",
-        "graphMode": "none",
-        "justifyMode": "auto",
-        "orientation": "auto",
-        "reduceOptions": {
-          "calcs": [
-            "lastNotNull"
-          ],
-          "fields": "",
-          "values": false
-        },
-        "textMode": "auto"
-      },
-      "pluginVersion": "12.3.1",
-      "targets": [
-        {
-          "editorMode": "code",
+          "range": true,
+          "refId": "A"
+        }
+      ],
+      "title": "Queue Backlog (Norm)",
+      "type": "stat"
+    },
+    {
+      "datasource": {
+        "type": "prometheus",
+        "uid": "PBFA97CFB590B2093"
+      },
+      "fieldConfig": {
+        "defaults": {
+          "color": {
+            "mode": "thresholds"
+          },
+          "decimals": 3,
+          "mappings": [],
+          "thresholds": {
+            "mode": "absolute",
+            "steps": [
+              {
+                "color": "green",
+                "value": null
+              },
+              {
+                "color": "orange",
+                "value": 0.4
+              },
+              {
+                "color": "red",
+                "value": 0.8
+              }
+            ]
+          },
+          "unit": "short"
+        },
+        "overrides": []
+      },
+      "gridPos": {
+        "h": 4,
+        "w": 6,
+        "x": 12,
+        "y": 0
+      },
+      "id": 3,
+      "options": {
+        "colorMode": "value",
+        "graphMode": "none",
+        "justifyMode": "auto",
+        "orientation": "auto",
+        "reduceOptions": {
+          "calcs": [
+            "lastNotNull"
+          ],
+          "fields": "",
+          "values": false
+        },
+        "textMode": "auto"
+      },
+      "pluginVersion": "12.3.1",
+      "targets": [
+        {
+          "editorMode": "code",
           "expr": "scalar(avg(last_over_time(antiatropos_average_latency_norm{mode=~\"live|simulated|hybrid|aws\"}[1m])))",
           "legendFormat": "latency (all modes)",
-          "range": true,
-          "refId": "A"
-        }
-      ],
-      "title": "Latency (Norm)",
-      "type": "stat"
-    },
-    {
-      "datasource": {
-        "type": "prometheus",
-        "uid": "PBFA97CFB590B2093"
-      },
-      "fieldConfig": {
-        "defaults": {
-          "color": {
-            "mode": "thresholds"
-          },
-          "decimals": 3,
-          "mappings": [],
-          "thresholds": {
-            "mode": "absolute",
-            "steps": [
-              {
-                "color": "green",
-                "value": null
-              },
-              {
-                "color": "orange",
-                "value": 100
-              },
-              {
-                "color": "red",
-                "value": 1000
-              }
-            ]
-          },
-          "unit": "short"
-        },
-        "overrides": []
-      },
-      "gridPos": {
-        "h": 4,
-        "w": 6,
-        "x": 18,
-        "y": 0
-      },
-      "id": 4,
-      "options": {
-        "colorMode": "value",
-        "graphMode": "none",
-        "justifyMode": "auto",
-        "orientation": "auto",
-        "reduceOptions": {
-          "calcs": [
-            "lastNotNull"
-          ],
-          "fields": "",
-          "values": false
-        },
-        "textMode": "auto"
-      },
-      "pluginVersion": "12.3.1",
-      "targets": [
-        {
-          "editorMode": "code",
+          "range": true,
+          "refId": "A"
+        }
+      ],
+      "title": "Latency (Norm)",
+      "type": "stat"
+    },
+    {
+      "datasource": {
+        "type": "prometheus",
+        "uid": "PBFA97CFB590B2093"
+      },
+      "fieldConfig": {
+        "defaults": {
+          "color": {
+            "mode": "thresholds"
+          },
+          "decimals": 3,
+          "mappings": [],
+          "thresholds": {
+            "mode": "absolute",
+            "steps": [
+              {
+                "color": "green",
+                "value": null
+              },
+              {
+                "color": "orange",
+                "value": 100
+              },
+              {
+                "color": "red",
+                "value": 1000
+              }
+            ]
+          },
+          "unit": "short"
+        },
+        "overrides": []
+      },
+      "gridPos": {
+        "h": 4,
+        "w": 6,
+        "x": 18,
+        "y": 0
+      },
+      "id": 4,
+      "options": {
+        "colorMode": "value",
+        "graphMode": "none",
+        "justifyMode": "auto",
+        "orientation": "auto",
+        "reduceOptions": {
+          "calcs": [
+            "lastNotNull"
+          ],
+          "fields": "",
+          "values": false
+        },
+        "textMode": "auto"
+      },
+      "pluginVersion": "12.3.1",
+      "targets": [
+        {
+          "editorMode": "code",
           "expr": "scalar(avg(last_over_time(antiatropos_lyapunov_energy{mode=~\"live|simulated|hybrid|aws\"}[1m])))",
           "legendFormat": "lyapunov energy (all modes)",
-          "range": true,
-          "refId": "A"
-        }
-      ],
-      "title": "Lyapunov Energy",
-      "type": "stat"
-    },
-    {
-      "datasource": {
-        "type": "prometheus",
-        "uid": "PBFA97CFB590B2093"
-      },
-      "fieldConfig": {
-        "defaults": {
-          "color": {
-            "mode": "palette-classic"
-          },
-          "custom": {
-            "axisBorderShow": false,
-            "axisCenteredZero": false,
-            "axisColorMode": "text",
-            "axisLabel": "",
-            "axisPlacement": "auto",
-            "barAlignment": 0,
-            "drawStyle": "line",
-            "fillOpacity": 10,
-            "gradientMode": "none",
-            "hideFrom": {
-              "legend": false,
-              "tooltip": false,
-              "viz": false
-            },
-            "insertNulls": false,
-            "lineInterpolation": "linear",
-            "lineWidth": 2,
-            "pointSize": 3,
-            "scaleDistribution": {
-              "type": "linear"
-            },
-            "showPoints": "auto",
-            "spanNulls": false,
-            "stacking": {
-              "group": "A",
-              "mode": "none"
-            },
-            "thresholdsStyle": {
-              "mode": "off"
-            }
-          },
-          "mappings": [],
-          "thresholds": {
-            "mode": "absolute",
-            "steps": [
-              {
-                "color": "green",
-                "value": null
-              },
-              {
-                "color": "red",
-                "value": 80
-              }
-            ]
-          },
-          "unit": "short"
-        },
-        "overrides": []
-      },
-      "gridPos": {
-        "h": 8,
-        "w": 12,
-        "x": 0,
-        "y": 4
-      },
-      "id": 5,
-      "options": {
-        "legend": {
-          "calcs": [],
-          "displayMode": "list",
-          "placement": "bottom",
-          "showLegend": true
-        },
-        "tooltip": {
-          "mode": "single",
-          "sort": "none"
-        }
-      },
-      "pluginVersion": "12.3.1",
-      "targets": [
-        {
-          "editorMode": "code",
+          "range": true,
+          "refId": "A"
+        }
+      ],
+      "title": "Lyapunov Energy",
+      "type": "stat"
+    },
+    {
+      "datasource": {
+        "type": "prometheus",
+        "uid": "PBFA97CFB590B2093"
+      },
+      "fieldConfig": {
+        "defaults": {
+          "color": {
+            "mode": "palette-classic"
+          },
+          "custom": {
+            "axisBorderShow": false,
+            "axisCenteredZero": false,
+            "axisColorMode": "text",
+            "axisLabel": "",
+            "axisPlacement": "auto",
+            "barAlignment": 0,
+            "drawStyle": "line",
+            "fillOpacity": 10,
+            "gradientMode": "none",
+            "hideFrom": {
+              "legend": false,
+              "tooltip": false,
+              "viz": false
+            },
+            "insertNulls": false,
+            "lineInterpolation": "linear",
+            "lineWidth": 2,
+            "pointSize": 3,
+            "scaleDistribution": {
+              "type": "linear"
+            },
+            "showPoints": "auto",
+            "spanNulls": false,
+            "stacking": {
+              "group": "A",
+              "mode": "none"
+            },
+            "thresholdsStyle": {
+              "mode": "off"
+            }
+          },
+          "mappings": [],
+          "thresholds": {
+            "mode": "absolute",
+            "steps": [
+              {
+                "color": "green",
+                "value": null
+              },
+              {
+                "color": "red",
+                "value": 80
+              }
+            ]
+          },
+          "unit": "short"
+        },
+        "overrides": []
+      },
+      "gridPos": {
+        "h": 8,
+        "w": 12,
+        "x": 0,
+        "y": 4
+      },
+      "id": 5,
+      "options": {
+        "legend": {
+          "calcs": [],
+          "displayMode": "list",
+          "placement": "bottom",
+          "showLegend": true
+        },
+        "tooltip": {
+          "mode": "single",
+          "sort": "none"
+        }
+      },
+      "pluginVersion": "12.3.1",
+      "targets": [
+        {
+          "editorMode": "code",
           "expr": "antiatropos_reward{mode=~\"live|simulated|hybrid|aws\"}",
           "legendFormat": "reward {{task_id}} ({{mode}})",
-          "range": true,
-          "refId": "A"
-        },
-        {
-          "editorMode": "code",
+          "range": true,
+          "refId": "A"
+        },
+        {
+          "editorMode": "code",
           "expr": "antiatropos_lyapunov_energy{mode=~\"live|simulated|hybrid|aws\"}",
           "legendFormat": "lyapunov {{task_id}} ({{mode}})",
-          "range": true,
-          "refId": "B"
-        }
-      ],
-      "title": "Reward and Lyapunov Trend",
-      "type": "timeseries"
-    },
-    {
-      "datasource": {
-        "type": "prometheus",
-        "uid": "PBFA97CFB590B2093"
-      },
-      "fieldConfig": {
-        "defaults": {
-          "color": {
-            "mode": "palette-classic"
-          },
-          "custom": {
-            "axisBorderShow": false,
-            "axisCenteredZero": false,
-            "axisColorMode": "text",
-            "axisLabel": "",
-            "axisPlacement": "auto",
-            "barAlignment": 0,
-            "drawStyle": "line",
-            "fillOpacity": 10,
-            "gradientMode": "none",
-            "hideFrom": {
-              "legend": false,
-              "tooltip": false,
-              "viz": false
-            },
-            "insertNulls": false,
-            "lineInterpolation": "linear",
-            "lineWidth": 2,
-            "pointSize": 3,
-            "scaleDistribution": {
-              "type": "linear"
-            },
-            "showPoints": "auto",
-            "spanNulls": false,
-            "stacking": {
-              "group": "A",
-              "mode": "none"
-            },
-            "thresholdsStyle": {
-              "mode": "off"
-            }
-          },
-          "mappings": [],
-          "thresholds": {
-            "mode": "absolute",
-            "steps": [
-              {
-                "color": "green",
-                "value": null
-              },
-              {
-                "color": "red",
-                "value": 80
-              }
-            ]
-          },
-          "unit": "short"
-        },
-        "overrides": []
-      },
-      "gridPos": {
-        "h": 8,
-        "w": 12,
-        "x": 12,
-        "y": 4
-      },
-      "id": 6,
-      "options": {
-        "legend": {
-          "calcs": [],
-          "displayMode": "list",
-          "placement": "bottom",
-          "showLegend": true
-        },
-        "tooltip": {
-          "mode": "single",
-          "sort": "none"
-        }
-      },
-      "pluginVersion": "12.3.1",
-      "targets": [
-        {
-          "editorMode": "code",
+          "range": true,
+          "refId": "B"
+        }
+      ],
+      "title": "Reward and Lyapunov Trend",
+      "type": "timeseries"
+    },
+    {
+      "datasource": {
+        "type": "prometheus",
+        "uid": "PBFA97CFB590B2093"
+      },
+      "fieldConfig": {
+        "defaults": {
+          "color": {
+            "mode": "palette-classic"
+          },
+          "custom": {
+            "axisBorderShow": false,
+            "axisCenteredZero": false,
+            "axisColorMode": "text",
+            "axisLabel": "",
+            "axisPlacement": "auto",
+            "barAlignment": 0,
+            "drawStyle": "line",
+            "fillOpacity": 10,
+            "gradientMode": "none",
+            "hideFrom": {
+              "legend": false,
+              "tooltip": false,
+              "viz": false
+            },
+            "insertNulls": false,
+            "lineInterpolation": "linear",
+            "lineWidth": 2,
+            "pointSize": 3,
+            "scaleDistribution": {
+              "type": "linear"
+            },
+            "showPoints": "auto",
+            "spanNulls": false,
+            "stacking": {
+              "group": "A",
+              "mode": "none"
+            },
+            "thresholdsStyle": {
+              "mode": "off"
+            }
+          },
+          "mappings": [],
+          "thresholds": {
+            "mode": "absolute",
+            "steps": [
+              {
+                "color": "green",
+                "value": null
+              },
+              {
+                "color": "red",
+                "value": 80
+              }
+            ]
+          },
+          "unit": "short"
+        },
+        "overrides": []
+      },
+      "gridPos": {
+        "h": 8,
+        "w": 12,
+        "x": 12,
+        "y": 4
+      },
+      "id": 6,
+      "options": {
+        "legend": {
+          "calcs": [],
+          "displayMode": "list",
+          "placement": "bottom",
+          "showLegend": true
+        },
+        "tooltip": {
+          "mode": "single",
+          "sort": "none"
+        }
+      },
+      "pluginVersion": "12.3.1",
+      "targets": [
+        {
+          "editorMode": "code",
           "expr": "antiatropos_total_queue_backlog{mode=~\"live|simulated|hybrid|aws\"}",
           "legendFormat": "queue {{task_id}} ({{mode}})",
-          "range": true,
-          "refId": "A"
-        },
-        {
-          "editorMode": "code",
+          "range": true,
+          "refId": "A"
+        },
+        {
+          "editorMode": "code",
           "expr": "antiatropos_average_latency_norm{mode=~\"live|simulated|hybrid|aws\"}",
           "legendFormat": "latency {{task_id}} ({{mode}})",
-          "range": true,
-          "refId": "B"
-        }
-      ],
-      "title": "Queue and Latency Trend",
-      "type": "timeseries"
-    },
-    {
-      "datasource": {
-        "type": "prometheus",
-        "uid": "PBFA97CFB590B2093"
-      },
-      "fieldConfig": {
-        "defaults": {
-          "color": {
-            "mode": "palette-classic"
-          },
-          "mappings": [],
-          "thresholds": {
-            "mode": "absolute",
-            "steps": [
-              {
-                "color": "green",
-                "value": null
-              },
-              {
-                "color": "red",
-                "value": 80
-              }
-            ]
-          },
-          "unit": "short"
-        },
-        "overrides": []
-      },
-      "gridPos": {
-        "h": 8,
-        "w": 12,
-        "x": 0,
-        "y": 12
-      },
-      "id": 7,
-      "options": {
-        "legend": {
-          "calcs": [],
-          "displayMode": "list",
-          "placement": "bottom",
-          "showLegend": true
-        },
-        "tooltip": {
-          "mode": "single",
-          "sort": "none"
-        }
-      },
-      "pluginVersion": "12.3.1",
-      "targets": [
-        {
-          "editorMode": "code",
+          "range": true,
+          "refId": "B"
+        }
+      ],
+      "title": "Queue and Latency Trend",
+      "type": "timeseries"
+    },
+    {
+      "datasource": {
+        "type": "prometheus",
+        "uid": "PBFA97CFB590B2093"
+      },
+      "fieldConfig": {
+        "defaults": {
+          "color": {
+            "mode": "palette-classic"
+          },
+          "mappings": [],
+          "thresholds": {
+            "mode": "absolute",
+            "steps": [
+              {
+                "color": "green",
+                "value": null
+              },
+              {
+                "color": "red",
+                "value": 80
+              }
+            ]
+          },
+          "unit": "short"
+        },
+        "overrides": []
+      },
+      "gridPos": {
+        "h": 8,
+        "w": 12,
+        "x": 0,
+        "y": 12
+      },
+      "id": 7,
+      "options": {
+        "legend": {
+          "calcs": [],
+          "displayMode": "list",
+          "placement": "bottom",
+          "showLegend": true
+        },
+        "tooltip": {
+          "mode": "single",
+          "sort": "none"
+        }
+      },
+      "pluginVersion": "12.3.1",
+      "targets": [
+        {
+          "editorMode": "code",
           "expr": "sum by (task_id, mode) (rate(antiatropos_steps_total{mode=~\"live|simulated|hybrid|aws\"}[1m]))",
           "legendFormat": "steps/sec {{task_id}} ({{mode}})",
-          "range": true,
-          "refId": "A"
-        },
-        {
-          "editorMode": "code",
+          "range": true,
+          "refId": "A"
+        },
+        {
+          "editorMode": "code",
           "expr": "sum by (task_id, mode, action_type) (rate(antiatropos_actions_total{mode=~\"live|simulated|hybrid|aws\"}[1m]))",
           "legendFormat": "actions/sec {{action_type}} ({{task_id}}, {{mode}})",
-          "range": true,
-          "refId": "B"
-        }
-      ],
-      "title": "Action Throughput",
-      "type": "timeseries"
-    },
-    {
-      "datasource": {
-        "type": "prometheus",
-        "uid": "PBFA97CFB590B2093"
-      },
-      "fieldConfig": {
-        "defaults": {
-          "color": {
-            "mode": "palette-classic"
-          },
-          "mappings": [],
-          "thresholds": {
-            "mode": "absolute",
-            "steps": [
-              {
-                "color": "green",
-                "value": null
-              },
-              {
-                "color": "red",
-                "value": 1
-              }
-            ]
-          },
-          "unit": "short"
-        },
-        "overrides": []
-      },
-      "gridPos": {
-        "h": 8,
-        "w": 12,
-        "x": 12,
-        "y": 12
-      },
-      "id": 8,
-      "options": {
-        "legend": {
-          "calcs": [],
-          "displayMode": "list",
-          "placement": "bottom",
-          "showLegend": true
-        },
-        "tooltip": {
-          "mode": "single",
-          "sort": "none"
-        }
-      },
-      "pluginVersion": "12.3.1",
-      "targets": [
-        {
-          "editorMode": "code",
+          "range": true,
+          "refId": "B"
+        }
+      ],
+      "title": "Action Throughput",
+      "type": "timeseries"
+    },
+    {
+      "datasource": {
+        "type": "prometheus",
+        "uid": "PBFA97CFB590B2093"
+      },
+      "fieldConfig": {
+        "defaults": {
+          "color": {
+            "mode": "palette-classic"
+          },
+          "mappings": [],
+          "thresholds": {
+            "mode": "absolute",
+            "steps": [
+              {
+                "color": "green",
+                "value": null
+              },
+              {
+                "color": "red",
+                "value": 1
+              }
+            ]
+          },
+          "unit": "short"
+        },
+        "overrides": []
+      },
+      "gridPos": {
+        "h": 8,
+        "w": 12,
+        "x": 12,
+        "y": 12
+      },
+      "id": 8,
+      "options": {
+        "legend": {
+          "calcs": [],
+          "displayMode": "list",
+          "placement": "bottom",
+          "showLegend": true
+        },
+        "tooltip": {
+          "mode": "single",
+          "sort": "none"
+        }
+      },
+      "pluginVersion": "12.3.1",
+      "targets": [
+        {
+          "editorMode": "code",
           "expr": "sum by (mode, error_code) (rate(antiatropos_executor_errors_total{mode=~\"live|simulated|hybrid|aws\"}[5m]))",
           "legendFormat": "executor errors {{error_code}} ({{mode}})",
-          "range": true,
-          "refId": "A"
-        },
-        {
-          "editorMode": "code",
+          "range": true,
+          "refId": "A"
+        },
+        {
+          "editorMode": "code",
           "expr": "histogram_quantile(0.95, sum(rate(antiatropos_executor_latency_ms_bucket{mode=~\"live|simulated|hybrid|aws\"}[5m])) by (le, mode))",
-          "legendFormat": "p95 executor latency {{mode}}",
-          "range": true,
-          "refId": "B"
-        }
-      ],
-      "title": "Executor Reliability",
-      "type": "timeseries"
-    }
-  ],
-  "refresh": "5s",
-  "schemaVersion": 41,
-  "style": "dark",
-  "tags": [
-    "antiatropos",
-    "sre"
-  ],
-  "templating": {
-    "list": []
-  },
-  "time": {
-    "from": "now-30m",
-    "to": "now"
-  },
-  "timepicker": {},
-  "timezone": "browser",
-  "title": "AntiAtropos Overview",
-  "uid": "antiatropos-overview",
+          "legendFormat": "p95 executor latency {{mode}}",
+          "range": true,
+          "refId": "B"
+        }
+      ],
+      "title": "Executor Reliability",
+      "type": "timeseries"
+    }
+  ],
+  "refresh": "5s",
+  "schemaVersion": 41,
+  "style": "dark",
+  "tags": [
+    "antiatropos",
+    "sre"
+  ],
+  "templating": {
+    "list": []
+  },
+  "time": {
+    "from": "now-30m",
+    "to": "now"
+  },
+  "timepicker": {},
+  "timezone": "browser",
+  "title": "AntiAtropos Overview",
+  "uid": "antiatropos-overview",
   "version": 2,
-  "weekStart": ""
-}
-
-
-
-
-
+  "weekStart": ""
+}
+
+
+
+
+

smoke_test.py

@@ -1,602 +1,602 @@
-#!/usr/bin/env python3
-"""
-AntiAtropos Local Smoke Test — 5-Node Validation.
-
-Validates simulator physics, reward signals, and grading WITHOUT any LLM,
-Colab, or AWS infrastructure. Uses only stdlib + project modules
-(simulator, stability, curriculum have zero external deps).
-
-Run from project root:
-    python smoke_test.py
-"""
-
-import sys
-import os
-import random
-import math
-
-# ── Make standalone imports work ──
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-
-from simulator import (
-    ClusterSimulator, NodeStatus, DEFAULT_CAPACITY, MAX_CAPACITY,
-    VIP_NODE_WEIGHTS, CRITICAL_NODES, COST_PER_CAPACITY_UNIT_PER_HOUR,
-    T1_INITIAL_LAMBDA, T2_INITIAL_LAMBDA, T3_INITIAL_LAMBDA,
-)
-from stability import (
-    compute_lyapunov, compute_reward, compute_barrier,
-    normalize_reward, smooth_sla_penalty, compute_drift,
-)
-from curriculum import CurriculumTracker, CURRICULUM
-
-# ── Test harness ─────────────────────────────────────────────────────────────────
-
-PASS = "PASS"
-FAIL = "FAIL"
-results: list[tuple[str, str, str]] = []  # (name, status, detail)
-
-
-def record(name: str, status: str, detail: str = "") -> None:
-    results.append((name, status, detail))
-    icon = "+" if status == PASS else "X"
-    msg = f"  [{icon}] {name}"
-    if detail:
-        msg += f" -- {detail}"
-    print(msg)
-
-
-def random_action(sim: ClusterSimulator) -> object:
-    """Generate a random valid action."""
-    node_ids = [n.node_id for n in sim._nodes]
-    action_types = ["SCALE_UP", "SCALE_DOWN", "REROUTE_TRAFFIC", "SHED_LOAD", "NO_OP"]
-
-    class _A:
-        pass
-
-    a = _A()
-    a.action_type = random.choice(action_types)
-    a.target_node_id = random.choice(node_ids)
-    a.parameter = round(random.random(), 2)
-    return a
-
-
-def run_episode(
-    sim: ClusterSimulator,
-    task_id: str,
-    max_steps: int = 60,
-    seed: int = 42,
-    action_policy: str = "random",
-) -> dict:
-    """
-    Run a full episode and collect diagnostics.
-
-    action_policy: 'random' | 'noop' | 'scale_up_vip'
-    """
-    sim.reset(task_id=task_id, seed=seed)
-
-    rewards_raw: list[float] = []
-    rewards_norm: list[float] = []
-    lyapunov_history: list[float] = []
-    sla_violations = 0
-    prev_v = 0.0
-    MAX_QUEUE_NORM = 200.0
-    MAX_LATENCY_NORM = 1000.0
-    ALPHA, BETA, GAMMA, DELTA = 0.002, 0.01, 10.0, 0.005
-
-    for step in range(1, max_steps + 1):
-        # Choose action
-        if action_policy == "noop":
-            class _A:
-                pass
-            a = _A()
-            a.action_type = "NO_OP"
-            a.target_node_id = "node-0"
-            a.parameter = 0.0
-        elif action_policy == "scale_up_vip":
-            class _A:
-                pass
-            a = _A()
-            a.action_type = "SCALE_UP"
-            a.target_node_id = "node-0"
-            a.parameter = 0.8
-        else:
-            a = random_action(sim)
-
-        sim.apply_action(a)
-        sim.tick()
-
-        # Compute reward (mirrors environment.py logic)
-        nodes_true = sim.state(for_agent=False)
-        current_v = compute_lyapunov(nodes_true)
-
-        # Avg latency (importance-weighted)
-        w_lat = 0.0
-        w_sum = 0.0
-        for n in nodes_true:
-            w = n.get("importance_weight", 1.0)
-            lat = MAX_LATENCY_NORM if n["status"] == NodeStatus.FAILED else n["latency_ms"]
-            w_lat += w * lat
-            w_sum += w
-        avg_lat_norm = min(1.0, max(0.0, (w_lat / w_sum / MAX_LATENCY_NORM) if w_sum > 0 else 1.0))
-
-        # Error rate
-        total_in = sum(
-            n.get("incoming_request_rate", 0) * n.get("importance_weight", 1.0)
-            for n in nodes_true
-        )
-        total_drop = sum(
-            n.get("dropped_requests", 0) * n.get("importance_weight", 1.0)
-            for n in nodes_true
-        )
-        error_rate = min(1.0, total_drop / total_in) if total_in > 0 else 0.0
-
-        sla_step = smooth_sla_penalty(avg_lat_norm, error_rate)
-        if avg_lat_norm > 0.20 or error_rate > 0.05:
-            sla_violations += 1
-
-        # Cost
-        total_cap = 0
-        for n in nodes_true:
-            if n["status"] != NodeStatus.FAILED:
-                total_cap += int(n.get("capacity_units", 0)) + int(n.get("pending_capacity_units", 0))
-        cost = total_cap * COST_PER_CAPACITY_UNIT_PER_HOUR
-
-        barrier = compute_barrier(nodes_true)
-        raw_r = compute_reward(
-            prev_v, current_v, cost, sla_step, ALPHA, BETA, GAMMA, barrier, DELTA
-        )
-        norm_r = normalize_reward(raw_r)
-
-        rewards_raw.append(raw_r)
-        rewards_norm.append(norm_r)
-        lyapunov_history.append(current_v)
-        prev_v = current_v
-
-    return {
-        "rewards_raw": rewards_raw,
-        "rewards_norm": rewards_norm,
-        "lyapunov_history": lyapunov_history,
-        "final_state": sim.state(for_agent=False),
-        "invalid_count": sim.invalid_action_count,
-        "sla_violations": sla_violations,
-    }
-
-
-# ════════════════════════════════════════════════════════════════════════════════
-# TEST FUNCTIONS
-# ════════════════════════════════════════════════════════════════════════════════
-
-def test_simulator_node_count():
-    """Simulator creates exactly 10 nodes; node-0 is VIP."""
-    print("\n--- Simulator Node Count ---")
-    sim = ClusterSimulator(n_nodes=5, task_id="task-1", seed=1)
-    nodes = sim.state(for_agent=False)
-
-    record("10 nodes created",
-           PASS if len(nodes) == 10 else FAIL,
-           f"got {len(nodes)}")
-
-    record("node-0 is VIP",
-           PASS if nodes[0]["is_vip"] else FAIL,
-           f"is_vip={nodes[0]['is_vip']}")
-
-    record("node-0 weight=4.0",
-           PASS if nodes[0]["importance_weight"] == 4.0 else FAIL,
-           f"weight={nodes[0]['importance_weight']}")
-
-    non_vip_weights = [n["importance_weight"] for n in nodes[1:]]
-    record("Non-VIP weight=1.0",
-           PASS if all(w == 1.0 for w in non_vip_weights) else FAIL,
-           f"unique weights={set(non_vip_weights)}")
-
-    node_ids = [n["node_id"] for n in nodes]
-    expected_ids = [f"node-{i}" for i in range(10)]
-    record("Node IDs 0-9",
-           PASS if node_ids == expected_ids else FAIL,
-           f"ids={node_ids}")
-
-    caps = [n["capacity_units"] for n in nodes]
-    record("All nodes at capacity 3",
-           PASS if all(c == 3 for c in caps) else FAIL,
-           f"caps={caps}")
-
-
-def test_task1_ramp():
-    """Task-1: traffic ramps, queues grow under NO_OP, rewards non-degenerate."""
-    print("\n--- Task-1: Linear Ramp (NO_OP policy) ---")
-    sim = ClusterSimulator(n_nodes=5, task_id="task-1")
-    ep = run_episode(sim, "task-1", max_steps=60, seed=42, action_policy="noop")
-
-    # Queues should grow (no scaling action taken)
-    final_queues = [n["queue_depth"] for n in ep["final_state"]]
-    max_q = max(final_queues)
-    record("Queues grow under NO_OP",
-           PASS if max_q > 0 else FAIL,
-           f"max_queue={max_q:.1f}")
-
-    # Rewards should not all be identical
-    unique_raw = len(set(round(r, 6) for r in ep["rewards_raw"]))
-    record("Raw rewards vary across steps",
-           PASS if unique_raw > 5 else FAIL,
-           f"unique values={unique_raw}/{len(ep['rewards_raw'])}")
-
-    # Normalized rewards in [0, 1]
-    all_in_range = all(0.0 <= r <= 1.0 for r in ep["rewards_norm"])
-    record("Normalized rewards in [0,1]",
-           PASS if all_in_range else FAIL,
-           f"min={min(ep['rewards_norm']):.4f} max={max(ep['rewards_norm']):.4f}")
-
-    # No NaN / inf
-    has_nan = any(math.isnan(r) or math.isinf(r) for r in ep["rewards_raw"])
-    record("No NaN/inf in raw rewards",
-           PASS if not has_nan else FAIL,
-           "")
-
-    # Lyapunov energy should trend upward (system destabilizing under NO_OP)
-    v_first5 = sum(ep["lyapunov_history"][:5]) / 5
-    v_last5 = sum(ep["lyapunov_history"][-5:]) / 5
-    record("Lyapunov energy rises under NO_OP",
-           PASS if v_last5 > v_first5 else FAIL,
-           f"early_avg={v_first5:.1f} late_avg={v_last5:.1f}")
-
-    print(f"  [i] SLA violations: {ep['sla_violations']}/60")
-    print(f"  [i] Avg norm reward: {sum(ep['rewards_norm'])/len(ep['rewards_norm']):.4f}")
-
-
-def test_task2_fault():
-    """Task-2: a node fails, queues react, reroute reduces load on failed node."""
-    print("\n--- Task-2: Fault Tolerance ---")
-    sim = ClusterSimulator(n_nodes=5, task_id="task-2")
-    ep = run_episode(sim, "task-2", max_steps=60, seed=42, action_policy="noop")
-
-    # At least one node should be FAILED by end (scripted failure)
-    failed = [n for n in ep["final_state"] if n["status"] == "FAILED"]
-    record("Scripted failure occurs",
-           PASS if len(failed) >= 1 else FAIL,
-           f"failed_nodes={len(failed)}")
-
-    # node-0 should NOT be the failed one (excluded from failure pool)
-    failed_ids = [n["node_id"] for n in failed]
-    record("node-0 not in failed set",
-           PASS if "node-0" not in failed_ids else FAIL,
-           f"failed_ids={failed_ids}")
-
-    # Rewards may plateau under NO_OP on constant-load tasks.
-    # Task-2 has fixed lambda, so steady-state reward has very low variance.
-    # This is expected — active policies (scale/reroute) create variation.
-    record("Raw rewards produced (may plateau under NO_OP)",
-           PASS if len(ep['rewards_raw']) == 60 else FAIL,
-           f"steps={len(ep['rewards_raw'])}")
-
-    # More importantly, normalized rewards should differ from 0.5 midpoint
-    # (proving the raw reward signal is non-trivial)
-    avg_norm = sum(ep['rewards_norm']) / len(ep['rewards_norm'])
-    record("Normalized reward is non-trivial (not stuck at 0.5)",
-           PASS if abs(avg_norm - 0.5) > 0.01 else FAIL,
-           f"avg_norm={avg_norm:.4f}")
-
-    # Normalized rewards in [0, 1]
-    all_in_range = all(0.0 <= r <= 1.0 for r in ep["rewards_norm"])
-    record("Normalized rewards in [0,1]",
-           PASS if all_in_range else FAIL,
-           f"min={min(ep['rewards_norm']):.4f} max={max(ep['rewards_norm']):.4f}")
-
-    # No NaN / inf
-    has_nan = any(math.isnan(r) or math.isinf(r) for r in ep["rewards_raw"])
-    record("No NaN/inf in raw rewards",
-           PASS if not has_nan else FAIL, "")
-
-    # Now test with targeted reroute on the scripted-failed node
-    # (NOT all nodes — rerouting everything to node-0 kills it)
-    sim2 = ClusterSimulator(n_nodes=5, task_id="task-2", seed=99)
-    sim2.reset(task_id="task-2", seed=99)
-    scripted_fail_id = None
-    for step in range(1, 61):
-        sim2.tick()
-        # Check if the scripted failure has been assigned
-        if sim2._failed_node_id and scripted_fail_id is None:
-            scripted_fail_id = sim2._failed_node_id
-            # Apply reroute specifically to the failed node
-            class _A:
-                pass
-            a = _A()
-            a.action_type = "REROUTE_TRAFFIC"
-            a.target_node_id = scripted_fail_id
-            a.parameter = 1.0
-            sim2.apply_action(a)
-            # Tick once more to see the effect
-            sim2.tick()
-            failed_node = next((n for n in sim2._nodes if n.node_id == scripted_fail_id), None)
-            base_share = sim2._t2_init_lambda / sim2._n_nodes
-            record("Reroute reduces failed node traffic",
-                   PASS if failed_node.incoming_request_rate < base_share else FAIL,
-                   f"node={scripted_fail_id} incoming={failed_node.incoming_request_rate:.1f} base_share={base_share:.1f}")
-            break
-
-
-def test_task3_surge():
-    """Task-3: surge hits node-1/node-2, SHED_LOAD on critical nodes rejected."""
-    print("\n--- Task-3: Periodic Surge ---")
-    sim = ClusterSimulator(n_nodes=5, task_id="task-3")
-    ep = run_episode(sim, "task-3", max_steps=60, seed=42, action_policy="noop")
-
-    # Rewards non-degenerate
-    unique_raw = len(set(round(r, 6) for r in ep["rewards_raw"]))
-    record("Raw rewards vary",
-           PASS if unique_raw > 5 else FAIL,
-           f"unique values={unique_raw}/{len(ep['rewards_raw'])}")
-
-    # Normalized rewards in [0, 1]
-    all_in_range = all(0.0 <= r <= 1.0 for r in ep["rewards_norm"])
-    record("Normalized rewards in [0,1]",
-           PASS if all_in_range else FAIL,
-           f"min={min(ep['rewards_norm']):.4f} max={max(ep['rewards_norm']):.4f}")
-
-    # No NaN / inf
-    has_nan = any(math.isnan(r) or math.isinf(r) for r in ep["rewards_raw"])
-    record("No NaN/inf in raw rewards",
-           PASS if not has_nan else FAIL, "")
-
-    # Test SHED_LOAD rejection on critical nodes
-    sim3 = ClusterSimulator(n_nodes=5, task_id="task-3", seed=7)
-    sim3.reset(task_id="task-3", seed=7)
-    for critical_id in CRITICAL_NODES:
-        class _A:
-            pass
-        a = _A()
-        a.action_type = "SHED_LOAD"
-        a.target_node_id = critical_id
-        a.parameter = 0.5
-        sim3.apply_action(a)
-    record("SHED_LOAD on critical nodes rejected",
-           PASS if sim3.invalid_action_count == len(CRITICAL_NODES) else FAIL,
-           f"invalid_count={sim3.invalid_action_count} expected={len(CRITICAL_NODES)}")
-
-    # SHED_LOAD on non-critical should be allowed
-    class _A2:
-        pass
-    a2 = _A2()
-    a2.action_type = "SHED_LOAD"
-    a2.target_node_id = "node-5"
-    a2.parameter = 0.5
-    sim3.apply_action(a2)
-    record("SHED_LOAD on non-critical node allowed",
-           PASS if sim3.invalid_action_count == len(CRITICAL_NODES) else FAIL,
-           f"invalid_count={sim3.invalid_action_count}")
-
-
-def test_scale_up_down():
-    """SCALE_UP increases capacity after boot delay; SCALE_DOWN decreases it."""
-    print("\n--- Scale Up / Scale Down ---")
-    sim = ClusterSimulator(n_nodes=5, task_id="task-1", seed=1)
-    sim.reset(task_id="task-1", seed=1)
-
-    # SCALE_UP node-3
-    class _A:
-        pass
-    a = _A()
-    a.action_type = "SCALE_UP"
-    a.target_node_id = "node-3"
-    a.parameter = 1.0  # 1 * MAX_SCALING_STEP=3 → 3 units
-    sim.apply_action(a)
-
-    # Check pending capacity before boot
-    node3 = next(n for n in sim._nodes if n.node_id == "node-3")
-    record("Pending capacity queued after SCALE_UP",
-           PASS if len(node3.pending_capacity_queue) > 0 else FAIL,
-           f"pending={len(node3.pending_capacity_queue)}")
-
-    # Tick through boot delay
-    for _ in range(6):
-        sim.tick()
-
-    node3 = next(n for n in sim._nodes if n.node_id == "node-3")
-    record("Capacity goes live after boot delay",
-           PASS if node3.capacity > DEFAULT_CAPACITY else FAIL,
-           f"capacity={node3.capacity}")
-
-    # SCALE_DOWN
-    prev_cap = node3.capacity
-    class _A2:
-        pass
-    a2 = _A2()
-    a2.action_type = "SCALE_DOWN"
-    a2.target_node_id = "node-3"
-    a2.parameter = 0.5
-    sim.apply_action(a2)
-    record("SCALE_DOWN reduces capacity",
-           PASS if node3.capacity < prev_cap else FAIL,
-           f"before={prev_cap} after={node3.capacity}")
-
-
-def test_reward_sanity():
-    """Detailed reward component sanity checks."""
-    print("\n--- Reward Sanity ---")
-
-    # Test normalize_reward mapping
-    r0 = normalize_reward(0.0)
-    record("normalize_reward(0.0) in [0,1]",
-           PASS if 0.0 <= r0 <= 1.0 else FAIL,
-           f"got {r0:.4f}")
-
-    r_neg = normalize_reward(-100.0)
-    r_pos = normalize_reward(100.0)
-    record("More negative raw -> lower normalized",
-           PASS if r_neg < r_pos else FAIL,
-           f"neg={r_neg:.4f} pos={r_pos:.4f}")
-
-    # Smooth SLA penalty
-    p_safe = smooth_sla_penalty(0.05, 0.01)   # well below thresholds
-    p_danger = smooth_sla_penalty(0.30, 0.10)  # above thresholds
-    record("SLA penalty: safe < danger",
-           PASS if p_safe < p_danger else FAIL,
-           f"safe={p_safe:.4f} danger={p_danger:.4f}")
-
-    # Barrier function
-    nodes_ok = [{"queue_depth": 50.0} for _ in range(10)]
-    nodes_bad = [{"queue_depth": 200.0} for _ in range(10)]
-    b_ok = compute_barrier(nodes_ok)
-    b_bad = compute_barrier(nodes_bad)
-    record("Barrier: safe queues < overloaded queues",
-           PASS if b_ok < b_bad else FAIL,
-           f"ok={b_ok:.1f} bad={b_bad:.1f}")
-    record("Barrier is 0 when all below Q_BARRIER_MAX",
-           PASS if b_ok == 0.0 else FAIL,
-           f"got {b_ok:.1f}")
-
-    # Lyapunov with VIP weight
-    nodes_no_vip = [{"queue_depth": 100.0, "importance_weight": 1.0} for _ in range(10)]
-    nodes_with_vip = [{"queue_depth": 100.0, "importance_weight": 4.0}] + \
-                     [{"queue_depth": 100.0, "importance_weight": 1.0} for _ in range(9)]
-    v_no_vip = compute_lyapunov(nodes_no_vip)
-    v_with_vip = compute_lyapunov(nodes_with_vip)
-    record("VIP weight amplifies Lyapunov energy",
-           PASS if v_with_vip > v_no_vip else FAIL,
-           f"no_vip={v_no_vip:.1f} with_vip={v_with_vip:.1f}")
-
-
-def test_grader_inline():
-    """Inline grader score validation (mirrors grader.py logic without importing it)."""
-    print("\n--- Grader Logic (Inline) ---")
-
-    # Baseline cost for 10 nodes at capacity 3
-    BASELINE = 10 * 3 * 0.05  # 1.50
-    MAX_COST = 10 * 5 * 0.05   # 2.50
-    COST_K = 3.0
-    TARGET_ENERGY = 2000.0
-    CURVE_POWER = 2.0
-
-    # Perfectly provisioned: cost == baseline -> score = 1.0
-    over_ratio = max(0.0, (BASELINE - BASELINE) / BASELINE)
-    cost_score = max(0.0, min(1.0, math.exp(-COST_K * over_ratio)))
-    record("Cost score=1.0 at baseline",
-           PASS if abs(cost_score - 1.0) < 1e-6 else FAIL,
-           f"got {cost_score:.4f}")
-
-    # 2x over-provisioned: score should be very low
-    over_ratio_2x = max(0.0, (2 * BASELINE - BASELINE) / BASELINE)
-    cost_score_2x = max(0.0, min(1.0, math.exp(-COST_K * over_ratio_2x)))
-    record("Cost score near 0 at 2x baseline",
-           PASS if cost_score_2x < 0.1 else FAIL,
-           f"got {cost_score_2x:.4f}")
-
-    # Stability: low energy -> high score
-    low_energy = 100.0
-    ratio = low_energy / TARGET_ENERGY
-    stab_score = 1.0 / (1.0 + ratio ** CURVE_POWER)
-    record("Stability score high at low energy",
-           PASS if stab_score > 0.9 else FAIL,
-           f"energy={low_energy} score={stab_score:.4f}")
-
-    # Stability: high energy -> low score
-    high_energy = 10000.0
-    ratio_h = high_energy / TARGET_ENERGY
-    stab_score_h = 1.0 / (1.0 + ratio_h ** CURVE_POWER)
-    record("Stability score low at high energy",
-           PASS if stab_score_h < 0.1 else FAIL,
-           f"energy={high_energy} score={stab_score_h:.4f}")
-
-
-def test_curriculum_tracker():
-    """Curriculum tracker advances stages on passing scores."""
-    print("\n--- Curriculum Tracker ---")
-    tracker = CurriculumTracker()
-
-    record("Starts at stage 0",
-           PASS if tracker.current_index == 0 else FAIL,
-           f"idx={tracker.current_index}")
-
-    record(f"Total stages = {len(CURRICULUM)}",
-           PASS if len(CURRICULUM) == 10 else FAIL,
-           f"got {len(CURRICULUM)}")
-
-    # Pass first stage
-    stage0 = tracker.current
-    passed = tracker.report_score(0.50)  # > 0.40 threshold
-    record("Pass stage 0 with score 0.50",
-           PASS if passed and tracker.current_index == 1 else FAIL,
-           f"passed={passed} idx={tracker.current_index}")
-
-    # Fail stage 1 (needs 0.50)
-    passed2 = tracker.report_score(0.30)  # < 0.50
-    record("Fail stage 1 with score 0.30",
-           PASS if not passed2 else FAIL,
-           f"passed={passed2} retries={tracker.current.retries}")
-
-    # Pass on retry
-    passed3 = tracker.report_score(0.60)
-    record("Pass stage 1 on retry with score 0.60",
-           PASS if passed3 and tracker.current_index == 2 else FAIL,
-           f"passed={passed3} idx={tracker.current_index}")
-
-    # Progress summary doesn't crash
-    summary = tracker.progress_summary()
-    record("progress_summary() returns string",
-           PASS if isinstance(summary, str) and len(summary) > 0 else FAIL,
-           f"len={len(summary)}")
-
-
-def test_cascade_and_recovery():
-    """Cascade failure detection and auto-recovery work."""
-    print("\n--- Cascade & Recovery ---")
-    sim = ClusterSimulator(n_nodes=5, task_id="task-1", seed=1)
-    sim.reset(task_id="task-1", seed=1)
-
-    # Artificially overload a node to trigger failure
-    node = sim._nodes[5]
-    node.queue_depth = 250.0  # > FATAL_FAIL_THRESHOLD=200
-    sim._update_statuses()
-    record("Node fails when queue > FATAL_FAIL_THRESHOLD",
-           PASS if node.status == NodeStatus.FAILED else FAIL,
-           f"status={node.status}")
-
-    record("Recovery timer set on overload failure",
-           PASS if node.recovery_timer > 0 else FAIL,
-           f"timer={node.recovery_timer}")
-
-    # Tick through recovery
-    for _ in range(25):
-        sim._process_recovery()
-
-    record("Node recovers after NODE_RECOVERY_TICKS",
-           PASS if node.status == NodeStatus.HEALTHY else FAIL,
-           f"status={node.status}")
-
-
-# ════════════════════════════════════════════════════════════════════════════════
-
-def main():
-    print("=" * 60)
-    print("AntiAtropos Smoke Test — 5-Node Cluster Validation")
-    print("=" * 60)
-
-    test_simulator_node_count()
-    test_task1_ramp()
-    test_task2_fault()
-    test_task3_surge()
-    test_scale_up_down()
-    test_reward_sanity()
-    test_grader_inline()
-    test_curriculum_tracker()
-    test_cascade_and_recovery()
-
-    # ── Summary ──
-    passed = sum(1 for _, s, _ in results if s == PASS)
-    failed = sum(1 for _, s, _ in results if s == FAIL)
-    total = len(results)
-
-    print("\n" + "=" * 60)
-    print(f"RESULTS: {passed}/{total} passed, {failed} failed")
-    print("=" * 60)
-
-    if failed > 0:
-        print("\nFailed tests:")
-        for name, status, detail in results:
-            if status == FAIL:
-                print(f"  X {name}: {detail}")
-
-    return 0 if failed == 0 else 1
-
-
-if __name__ == "__main__":
-    sys.exit(main())
+#!/usr/bin/env python3
+"""
+AntiAtropos Local Smoke Test — 5-Node Validation.
+
+Validates simulator physics, reward signals, and grading WITHOUT any LLM,
+Colab, or AWS infrastructure. Uses only stdlib + project modules
+(simulator, stability, curriculum have zero external deps).
+
+Run from project root:
+    python smoke_test.py
+"""
+
+import sys
+import os
+import random
+import math
+
+# ── Make standalone imports work ──
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+from simulator import (
+    ClusterSimulator, NodeStatus, DEFAULT_CAPACITY, MAX_CAPACITY,
+    VIP_NODE_WEIGHTS, CRITICAL_NODES, COST_PER_CAPACITY_UNIT_PER_HOUR,
+    T1_INITIAL_LAMBDA, T2_INITIAL_LAMBDA, T3_INITIAL_LAMBDA,
+)
+from stability import (
+    compute_lyapunov, compute_reward, compute_barrier,
+    normalize_reward, smooth_sla_penalty, compute_drift,
+)
+from curriculum import CurriculumTracker, CURRICULUM
+
+# ── Test harness ─────────────────────────────────────────────────────────────────
+
+PASS = "PASS"
+FAIL = "FAIL"
+results: list[tuple[str, str, str]] = []  # (name, status, detail)
+
+
+def record(name: str, status: str, detail: str = "") -> None:
+    results.append((name, status, detail))
+    icon = "+" if status == PASS else "X"
+    msg = f"  [{icon}] {name}"
+    if detail:
+        msg += f" -- {detail}"
+    print(msg)
+
+
+def random_action(sim: ClusterSimulator) -> object:
+    """Generate a random valid action."""
+    node_ids = [n.node_id for n in sim._nodes]
+    action_types = ["SCALE_UP", "SCALE_DOWN", "REROUTE_TRAFFIC", "SHED_LOAD", "NO_OP"]
+
+    class _A:
+        pass
+
+    a = _A()
+    a.action_type = random.choice(action_types)
+    a.target_node_id = random.choice(node_ids)
+    a.parameter = round(random.random(), 2)
+    return a
+
+
+def run_episode(
+    sim: ClusterSimulator,
+    task_id: str,
+    max_steps: int = 60,
+    seed: int = 42,
+    action_policy: str = "random",
+) -> dict:
+    """
+    Run a full episode and collect diagnostics.
+
+    action_policy: 'random' | 'noop' | 'scale_up_vip'
+    """
+    sim.reset(task_id=task_id, seed=seed)
+
+    rewards_raw: list[float] = []
+    rewards_norm: list[float] = []
+    lyapunov_history: list[float] = []
+    sla_violations = 0
+    prev_v = 0.0
+    MAX_QUEUE_NORM = 200.0
+    MAX_LATENCY_NORM = 1000.0
+    ALPHA, BETA, GAMMA, DELTA = 0.002, 0.01, 10.0, 0.005
+
+    for step in range(1, max_steps + 1):
+        # Choose action
+        if action_policy == "noop":
+            class _A:
+                pass
+            a = _A()
+            a.action_type = "NO_OP"
+            a.target_node_id = "node-0"
+            a.parameter = 0.0
+        elif action_policy == "scale_up_vip":
+            class _A:
+                pass
+            a = _A()
+            a.action_type = "SCALE_UP"
+            a.target_node_id = "node-0"
+            a.parameter = 0.8
+        else:
+            a = random_action(sim)
+
+        sim.apply_action(a)
+        sim.tick()
+
+        # Compute reward (mirrors environment.py logic)
+        nodes_true = sim.state(for_agent=False)
+        current_v = compute_lyapunov(nodes_true)
+
+        # Avg latency (importance-weighted)
+        w_lat = 0.0
+        w_sum = 0.0
+        for n in nodes_true:
+            w = n.get("importance_weight", 1.0)
+            lat = MAX_LATENCY_NORM if n["status"] == NodeStatus.FAILED else n["latency_ms"]
+            w_lat += w * lat
+            w_sum += w
+        avg_lat_norm = min(1.0, max(0.0, (w_lat / w_sum / MAX_LATENCY_NORM) if w_sum > 0 else 1.0))
+
+        # Error rate
+        total_in = sum(
+            n.get("incoming_request_rate", 0) * n.get("importance_weight", 1.0)
+            for n in nodes_true
+        )
+        total_drop = sum(
+            n.get("dropped_requests", 0) * n.get("importance_weight", 1.0)
+            for n in nodes_true
+        )
+        error_rate = min(1.0, total_drop / total_in) if total_in > 0 else 0.0
+
+        sla_step = smooth_sla_penalty(avg_lat_norm, error_rate)
+        if avg_lat_norm > 0.20 or error_rate > 0.05:
+            sla_violations += 1
+
+        # Cost
+        total_cap = 0
+        for n in nodes_true:
+            if n["status"] != NodeStatus.FAILED:
+                total_cap += int(n.get("capacity_units", 0)) + int(n.get("pending_capacity_units", 0))
+        cost = total_cap * COST_PER_CAPACITY_UNIT_PER_HOUR
+
+        barrier = compute_barrier(nodes_true)
+        raw_r = compute_reward(
+            prev_v, current_v, cost, sla_step, ALPHA, BETA, GAMMA, barrier, DELTA
+        )
+        norm_r = normalize_reward(raw_r)
+
+        rewards_raw.append(raw_r)
+        rewards_norm.append(norm_r)
+        lyapunov_history.append(current_v)
+        prev_v = current_v
+
+    return {
+        "rewards_raw": rewards_raw,
+        "rewards_norm": rewards_norm,
+        "lyapunov_history": lyapunov_history,
+        "final_state": sim.state(for_agent=False),
+        "invalid_count": sim.invalid_action_count,
+        "sla_violations": sla_violations,
+    }
+
+
+# ════════════════════════════════════════════════════════════════════════════════
+# TEST FUNCTIONS
+# ════════════════════════════════════════════════════════════════════════════════
+
+def test_simulator_node_count():
+    """Simulator creates exactly 10 nodes; node-0 is VIP."""
+    print("\n--- Simulator Node Count ---")
+    sim = ClusterSimulator(n_nodes=5, task_id="task-1", seed=1)
+    nodes = sim.state(for_agent=False)
+
+    record("10 nodes created",
+           PASS if len(nodes) == 10 else FAIL,
+           f"got {len(nodes)}")
+
+    record("node-0 is VIP",
+           PASS if nodes[0]["is_vip"] else FAIL,
+           f"is_vip={nodes[0]['is_vip']}")
+
+    record("node-0 weight=4.0",
+           PASS if nodes[0]["importance_weight"] == 4.0 else FAIL,
+           f"weight={nodes[0]['importance_weight']}")
+
+    non_vip_weights = [n["importance_weight"] for n in nodes[1:]]
+    record("Non-VIP weight=1.0",
+           PASS if all(w == 1.0 for w in non_vip_weights) else FAIL,
+           f"unique weights={set(non_vip_weights)}")
+
+    node_ids = [n["node_id"] for n in nodes]
+    expected_ids = [f"node-{i}" for i in range(10)]
+    record("Node IDs 0-9",
+           PASS if node_ids == expected_ids else FAIL,
+           f"ids={node_ids}")
+
+    caps = [n["capacity_units"] for n in nodes]
+    record("All nodes at capacity 3",
+           PASS if all(c == 3 for c in caps) else FAIL,
+           f"caps={caps}")
+
+
+def test_task1_ramp():
+    """Task-1: traffic ramps, queues grow under NO_OP, rewards non-degenerate."""
+    print("\n--- Task-1: Linear Ramp (NO_OP policy) ---")
+    sim = ClusterSimulator(n_nodes=5, task_id="task-1")
+    ep = run_episode(sim, "task-1", max_steps=60, seed=42, action_policy="noop")
+
+    # Queues should grow (no scaling action taken)
+    final_queues = [n["queue_depth"] for n in ep["final_state"]]
+    max_q = max(final_queues)
+    record("Queues grow under NO_OP",
+           PASS if max_q > 0 else FAIL,
+           f"max_queue={max_q:.1f}")
+
+    # Rewards should not all be identical
+    unique_raw = len(set(round(r, 6) for r in ep["rewards_raw"]))
+    record("Raw rewards vary across steps",
+           PASS if unique_raw > 5 else FAIL,
+           f"unique values={unique_raw}/{len(ep['rewards_raw'])}")
+
+    # Normalized rewards in [0, 1]
+    all_in_range = all(0.0 <= r <= 1.0 for r in ep["rewards_norm"])
+    record("Normalized rewards in [0,1]",
+           PASS if all_in_range else FAIL,
+           f"min={min(ep['rewards_norm']):.4f} max={max(ep['rewards_norm']):.4f}")
+
+    # No NaN / inf
+    has_nan = any(math.isnan(r) or math.isinf(r) for r in ep["rewards_raw"])
+    record("No NaN/inf in raw rewards",
+           PASS if not has_nan else FAIL,
+           "")
+
+    # Lyapunov energy should trend upward (system destabilizing under NO_OP)
+    v_first5 = sum(ep["lyapunov_history"][:5]) / 5
+    v_last5 = sum(ep["lyapunov_history"][-5:]) / 5
+    record("Lyapunov energy rises under NO_OP",
+           PASS if v_last5 > v_first5 else FAIL,
+           f"early_avg={v_first5:.1f} late_avg={v_last5:.1f}")
+
+    print(f"  [i] SLA violations: {ep['sla_violations']}/60")
+    print(f"  [i] Avg norm reward: {sum(ep['rewards_norm'])/len(ep['rewards_norm']):.4f}")
+
+
+def test_task2_fault():
+    """Task-2: a node fails, queues react, reroute reduces load on failed node."""
+    print("\n--- Task-2: Fault Tolerance ---")
+    sim = ClusterSimulator(n_nodes=5, task_id="task-2")
+    ep = run_episode(sim, "task-2", max_steps=60, seed=42, action_policy="noop")
+
+    # At least one node should be FAILED by end (scripted failure)
+    failed = [n for n in ep["final_state"] if n["status"] == "FAILED"]
+    record("Scripted failure occurs",
+           PASS if len(failed) >= 1 else FAIL,
+           f"failed_nodes={len(failed)}")
+
+    # node-0 should NOT be the failed one (excluded from failure pool)
+    failed_ids = [n["node_id"] for n in failed]
+    record("node-0 not in failed set",
+           PASS if "node-0" not in failed_ids else FAIL,
+           f"failed_ids={failed_ids}")
+
+    # Rewards may plateau under NO_OP on constant-load tasks.
+    # Task-2 has fixed lambda, so steady-state reward has very low variance.
+    # This is expected — active policies (scale/reroute) create variation.
+    record("Raw rewards produced (may plateau under NO_OP)",
+           PASS if len(ep['rewards_raw']) == 60 else FAIL,
+           f"steps={len(ep['rewards_raw'])}")
+
+    # More importantly, normalized rewards should differ from 0.5 midpoint
+    # (proving the raw reward signal is non-trivial)
+    avg_norm = sum(ep['rewards_norm']) / len(ep['rewards_norm'])
+    record("Normalized reward is non-trivial (not stuck at 0.5)",
+           PASS if abs(avg_norm - 0.5) > 0.01 else FAIL,
+           f"avg_norm={avg_norm:.4f}")
+
+    # Normalized rewards in [0, 1]
+    all_in_range = all(0.0 <= r <= 1.0 for r in ep["rewards_norm"])
+    record("Normalized rewards in [0,1]",
+           PASS if all_in_range else FAIL,
+           f"min={min(ep['rewards_norm']):.4f} max={max(ep['rewards_norm']):.4f}")
+
+    # No NaN / inf
+    has_nan = any(math.isnan(r) or math.isinf(r) for r in ep["rewards_raw"])
+    record("No NaN/inf in raw rewards",
+           PASS if not has_nan else FAIL, "")
+
+    # Now test with targeted reroute on the scripted-failed node
+    # (NOT all nodes — rerouting everything to node-0 kills it)
+    sim2 = ClusterSimulator(n_nodes=5, task_id="task-2", seed=99)
+    sim2.reset(task_id="task-2", seed=99)
+    scripted_fail_id = None
+    for step in range(1, 61):
+        sim2.tick()
+        # Check if the scripted failure has been assigned
+        if sim2._failed_node_id and scripted_fail_id is None:
+            scripted_fail_id = sim2._failed_node_id
+            # Apply reroute specifically to the failed node
+            class _A:
+                pass
+            a = _A()
+            a.action_type = "REROUTE_TRAFFIC"
+            a.target_node_id = scripted_fail_id
+            a.parameter = 1.0
+            sim2.apply_action(a)
+            # Tick once more to see the effect
+            sim2.tick()
+            failed_node = next((n for n in sim2._nodes if n.node_id == scripted_fail_id), None)
+            base_share = sim2._t2_init_lambda / sim2._n_nodes
+            record("Reroute reduces failed node traffic",
+                   PASS if failed_node.incoming_request_rate < base_share else FAIL,
+                   f"node={scripted_fail_id} incoming={failed_node.incoming_request_rate:.1f} base_share={base_share:.1f}")
+            break
+
+
+def test_task3_surge():
+    """Task-3: surge hits node-1/node-2, SHED_LOAD on critical nodes rejected."""
+    print("\n--- Task-3: Periodic Surge ---")
+    sim = ClusterSimulator(n_nodes=5, task_id="task-3")
+    ep = run_episode(sim, "task-3", max_steps=60, seed=42, action_policy="noop")
+
+    # Rewards non-degenerate
+    unique_raw = len(set(round(r, 6) for r in ep["rewards_raw"]))
+    record("Raw rewards vary",
+           PASS if unique_raw > 5 else FAIL,
+           f"unique values={unique_raw}/{len(ep['rewards_raw'])}")
+
+    # Normalized rewards in [0, 1]
+    all_in_range = all(0.0 <= r <= 1.0 for r in ep["rewards_norm"])
+    record("Normalized rewards in [0,1]",
+           PASS if all_in_range else FAIL,
+           f"min={min(ep['rewards_norm']):.4f} max={max(ep['rewards_norm']):.4f}")
+
+    # No NaN / inf
+    has_nan = any(math.isnan(r) or math.isinf(r) for r in ep["rewards_raw"])
+    record("No NaN/inf in raw rewards",
+           PASS if not has_nan else FAIL, "")
+
+    # Test SHED_LOAD rejection on critical nodes
+    sim3 = ClusterSimulator(n_nodes=5, task_id="task-3", seed=7)
+    sim3.reset(task_id="task-3", seed=7)
+    for critical_id in CRITICAL_NODES:
+        class _A:
+            pass
+        a = _A()
+        a.action_type = "SHED_LOAD"
+        a.target_node_id = critical_id
+        a.parameter = 0.5
+        sim3.apply_action(a)
+    record("SHED_LOAD on critical nodes rejected",
+           PASS if sim3.invalid_action_count == len(CRITICAL_NODES) else FAIL,
+           f"invalid_count={sim3.invalid_action_count} expected={len(CRITICAL_NODES)}")
+
+    # SHED_LOAD on non-critical should be allowed
+    class _A2:
+        pass
+    a2 = _A2()
+    a2.action_type = "SHED_LOAD"
+    a2.target_node_id = "node-5"
+    a2.parameter = 0.5
+    sim3.apply_action(a2)
+    record("SHED_LOAD on non-critical node allowed",
+           PASS if sim3.invalid_action_count == len(CRITICAL_NODES) else FAIL,
+           f"invalid_count={sim3.invalid_action_count}")
+
+
+def test_scale_up_down():
+    """SCALE_UP increases capacity after boot delay; SCALE_DOWN decreases it."""
+    print("\n--- Scale Up / Scale Down ---")
+    sim = ClusterSimulator(n_nodes=5, task_id="task-1", seed=1)
+    sim.reset(task_id="task-1", seed=1)
+
+    # SCALE_UP node-3
+    class _A:
+        pass
+    a = _A()
+    a.action_type = "SCALE_UP"
+    a.target_node_id = "node-3"
+    a.parameter = 1.0  # 1 * MAX_SCALING_STEP=3 → 3 units
+    sim.apply_action(a)
+
+    # Check pending capacity before boot
+    node3 = next(n for n in sim._nodes if n.node_id == "node-3")
+    record("Pending capacity queued after SCALE_UP",
+           PASS if len(node3.pending_capacity_queue) > 0 else FAIL,
+           f"pending={len(node3.pending_capacity_queue)}")
+
+    # Tick through boot delay
+    for _ in range(6):
+        sim.tick()
+
+    node3 = next(n for n in sim._nodes if n.node_id == "node-3")
+    record("Capacity goes live after boot delay",
+           PASS if node3.capacity > DEFAULT_CAPACITY else FAIL,
+           f"capacity={node3.capacity}")
+
+    # SCALE_DOWN
+    prev_cap = node3.capacity
+    class _A2:
+        pass
+    a2 = _A2()
+    a2.action_type = "SCALE_DOWN"
+    a2.target_node_id = "node-3"
+    a2.parameter = 0.5
+    sim.apply_action(a2)
+    record("SCALE_DOWN reduces capacity",
+           PASS if node3.capacity < prev_cap else FAIL,
+           f"before={prev_cap} after={node3.capacity}")
+
+
+def test_reward_sanity():
+    """Detailed reward component sanity checks."""
+    print("\n--- Reward Sanity ---")
+
+    # Test normalize_reward mapping
+    r0 = normalize_reward(0.0)
+    record("normalize_reward(0.0) in [0,1]",
+           PASS if 0.0 <= r0 <= 1.0 else FAIL,
+           f"got {r0:.4f}")
+
+    r_neg = normalize_reward(-100.0)
+    r_pos = normalize_reward(100.0)
+    record("More negative raw -> lower normalized",
+           PASS if r_neg < r_pos else FAIL,
+           f"neg={r_neg:.4f} pos={r_pos:.4f}")
+
+    # Smooth SLA penalty
+    p_safe = smooth_sla_penalty(0.05, 0.01)   # well below thresholds
+    p_danger = smooth_sla_penalty(0.30, 0.10)  # above thresholds
+    record("SLA penalty: safe < danger",
+           PASS if p_safe < p_danger else FAIL,
+           f"safe={p_safe:.4f} danger={p_danger:.4f}")
+
+    # Barrier function
+    nodes_ok = [{"queue_depth": 50.0} for _ in range(10)]
+    nodes_bad = [{"queue_depth": 200.0} for _ in range(10)]
+    b_ok = compute_barrier(nodes_ok)
+    b_bad = compute_barrier(nodes_bad)
+    record("Barrier: safe queues < overloaded queues",
+           PASS if b_ok < b_bad else FAIL,
+           f"ok={b_ok:.1f} bad={b_bad:.1f}")
+    record("Barrier is 0 when all below Q_BARRIER_MAX",
+           PASS if b_ok == 0.0 else FAIL,
+           f"got {b_ok:.1f}")
+
+    # Lyapunov with VIP weight
+    nodes_no_vip = [{"queue_depth": 100.0, "importance_weight": 1.0} for _ in range(10)]
+    nodes_with_vip = [{"queue_depth": 100.0, "importance_weight": 4.0}] + \
+                     [{"queue_depth": 100.0, "importance_weight": 1.0} for _ in range(9)]
+    v_no_vip = compute_lyapunov(nodes_no_vip)
+    v_with_vip = compute_lyapunov(nodes_with_vip)
+    record("VIP weight amplifies Lyapunov energy",
+           PASS if v_with_vip > v_no_vip else FAIL,
+           f"no_vip={v_no_vip:.1f} with_vip={v_with_vip:.1f}")
+
+
+def test_grader_inline():
+    """Inline grader score validation (mirrors grader.py logic without importing it)."""
+    print("\n--- Grader Logic (Inline) ---")
+
+    # Baseline cost for 10 nodes at capacity 3
+    BASELINE = 10 * 3 * 0.05  # 1.50
+    MAX_COST = 10 * 5 * 0.05   # 2.50
+    COST_K = 3.0
+    TARGET_ENERGY = 2000.0
+    CURVE_POWER = 2.0
+
+    # Perfectly provisioned: cost == baseline -> score = 1.0
+    over_ratio = max(0.0, (BASELINE - BASELINE) / BASELINE)
+    cost_score = max(0.0, min(1.0, math.exp(-COST_K * over_ratio)))
+    record("Cost score=1.0 at baseline",
+           PASS if abs(cost_score - 1.0) < 1e-6 else FAIL,
+           f"got {cost_score:.4f}")
+
+    # 2x over-provisioned: score should be very low
+    over_ratio_2x = max(0.0, (2 * BASELINE - BASELINE) / BASELINE)
+    cost_score_2x = max(0.0, min(1.0, math.exp(-COST_K * over_ratio_2x)))
+    record("Cost score near 0 at 2x baseline",
+           PASS if cost_score_2x < 0.1 else FAIL,
+           f"got {cost_score_2x:.4f}")
+
+    # Stability: low energy -> high score
+    low_energy = 100.0
+    ratio = low_energy / TARGET_ENERGY
+    stab_score = 1.0 / (1.0 + ratio ** CURVE_POWER)
+    record("Stability score high at low energy",
+           PASS if stab_score > 0.9 else FAIL,
+           f"energy={low_energy} score={stab_score:.4f}")
+
+    # Stability: high energy -> low score
+    high_energy = 10000.0
+    ratio_h = high_energy / TARGET_ENERGY
+    stab_score_h = 1.0 / (1.0 + ratio_h ** CURVE_POWER)
+    record("Stability score low at high energy",
+           PASS if stab_score_h < 0.1 else FAIL,
+           f"energy={high_energy} score={stab_score_h:.4f}")
+
+
+def test_curriculum_tracker():
+    """Curriculum tracker advances stages on passing scores."""
+    print("\n--- Curriculum Tracker ---")
+    tracker = CurriculumTracker()
+
+    record("Starts at stage 0",
+           PASS if tracker.current_index == 0 else FAIL,
+           f"idx={tracker.current_index}")
+
+    record(f"Total stages = {len(CURRICULUM)}",
+           PASS if len(CURRICULUM) == 10 else FAIL,
+           f"got {len(CURRICULUM)}")
+
+    # Pass first stage
+    stage0 = tracker.current
+    passed = tracker.report_score(0.50)  # > 0.40 threshold
+    record("Pass stage 0 with score 0.50",
+           PASS if passed and tracker.current_index == 1 else FAIL,
+           f"passed={passed} idx={tracker.current_index}")
+
+    # Fail stage 1 (needs 0.50)
+    passed2 = tracker.report_score(0.30)  # < 0.50
+    record("Fail stage 1 with score 0.30",
+           PASS if not passed2 else FAIL,
+           f"passed={passed2} retries={tracker.current.retries}")
+
+    # Pass on retry
+    passed3 = tracker.report_score(0.60)
+    record("Pass stage 1 on retry with score 0.60",
+           PASS if passed3 and tracker.current_index == 2 else FAIL,
+           f"passed={passed3} idx={tracker.current_index}")
+
+    # Progress summary doesn't crash
+    summary = tracker.progress_summary()
+    record("progress_summary() returns string",
+           PASS if isinstance(summary, str) and len(summary) > 0 else FAIL,
+           f"len={len(summary)}")
+
+
+def test_cascade_and_recovery():
+    """Cascade failure detection and auto-recovery work."""
+    print("\n--- Cascade & Recovery ---")
+    sim = ClusterSimulator(n_nodes=5, task_id="task-1", seed=1)
+    sim.reset(task_id="task-1", seed=1)
+
+    # Artificially overload a node to trigger failure
+    node = sim._nodes[5]
+    node.queue_depth = 250.0  # > FATAL_FAIL_THRESHOLD=200
+    sim._update_statuses()
+    record("Node fails when queue > FATAL_FAIL_THRESHOLD",
+           PASS if node.status == NodeStatus.FAILED else FAIL,
+           f"status={node.status}")
+
+    record("Recovery timer set on overload failure",
+           PASS if node.recovery_timer > 0 else FAIL,
+           f"timer={node.recovery_timer}")
+
+    # Tick through recovery
+    for _ in range(25):
+        sim._process_recovery()
+
+    record("Node recovers after NODE_RECOVERY_TICKS",
+           PASS if node.status == NodeStatus.HEALTHY else FAIL,
+           f"status={node.status}")
+
+
+# ════════════════════════════════════════════════════════════════════════════════
+
+def main():
+    print("=" * 60)
+    print("AntiAtropos Smoke Test — 5-Node Cluster Validation")
+    print("=" * 60)
+
+    test_simulator_node_count()
+    test_task1_ramp()
+    test_task2_fault()
+    test_task3_surge()
+    test_scale_up_down()
+    test_reward_sanity()
+    test_grader_inline()
+    test_curriculum_tracker()
+    test_cascade_and_recovery()
+
+    # ── Summary ──
+    passed = sum(1 for _, s, _ in results if s == PASS)
+    failed = sum(1 for _, s, _ in results if s == FAIL)
+    total = len(results)
+
+    print("\n" + "=" * 60)
+    print(f"RESULTS: {passed}/{total} passed, {failed} failed")
+    print("=" * 60)
+
+    if failed > 0:
+        print("\nFailed tests:")
+        for name, status, detail in results:
+            if status == FAIL:
+                print(f"  X {name}: {detail}")
+
+    return 0 if failed == 0 else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

train_test.py

@@ -1,538 +1,538 @@
-#!/usr/bin/env python3
-"""
-AntiAtropos Training Validation — Local Test.
-
-Validates the training pipeline (loss functions, episode collection, trainer)
-using a MockPolicyModel (no GPU needed). Run before going to Colab.
-
-Run from project root:
-    python train_test.py
-"""
-
-import sys
-import os
-import math
-import random
-
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-
-from training.losses import (
-    compute_returns, compute_gae,
-    reinforce_loss, reinforce_baseline_loss,
-    grpo_loss, rloo_loss,
-    normalize_rewards, compute_reward_stats,
-)
-from training.trainer import (
-    SRETrainer, TrainingConfig, EpisodeCollector,
-    MockPolicyModel,
-    LOSS_REINFORCE, LOSS_REINFORCE_BASELINE, LOSS_GRPO, LOSS_RLOO,
-)
-
-PASS = "PASS"
-FAIL = "FAIL"
-results: list[tuple[str, str, str]] = []
-
-
-def record(name: str, status: str, detail: str = "") -> None:
-    results.append((name, status, detail))
-    icon = "+" if status == PASS else "X"
-    msg = f"  [{icon}] {name}"
-    if detail:
-        msg += f" -- {detail}"
-    print(msg)
-
-
-# ════════════════════════════════════════════════════════════════════════════════
-# 1. Return Computation
-# ════════════════════════════════════════════════════════════════════════════════
-
-def test_returns():
-    print("\n--- Return Computation ---")
-    # Simple case: [1, 1, 1] with gamma=0.99
-    returns = compute_returns([1.0, 1.0, 1.0], gamma=0.99)
-    # G_2 = 1.0, G_1 = 1 + 0.99*1 = 1.99, G_0 = 1 + 0.99*1.99 = 2.9701
-    record("compute_returns[0]",
-           PASS if abs(returns[0] - 2.9701) < 0.001 else FAIL,
-           f"got {returns[0]:.4f} expected 2.9701")
-    record("compute_returns[2]",
-           PASS if abs(returns[2] - 1.0) < 0.001 else FAIL,
-           f"got {returns[2]:.4f} expected 1.0")
-
-    # Empty rewards
-    returns_empty = compute_returns([])
-    record("compute_returns handles empty",
-           PASS if returns_empty == [] else FAIL,
-           f"got {returns_empty}")
-
-    # Single reward
-    returns_single = compute_returns([5.0])
-    record("compute_returns single reward",
-           PASS if abs(returns_single[0] - 5.0) < 0.001 else FAIL,
-           f"got {returns_single[0]:.4f}")
-
-    # Discount factor = 0 → only immediate reward matters
-    returns_0 = compute_returns([1.0, 2.0, 3.0], gamma=0.0)
-    record("gamma=0: returns = rewards",
-           PASS if returns_0 == [1.0, 2.0, 3.0] else FAIL,
-           f"got {returns_0}")
-
-
-# ════════════════════════════════════════════════════════════════════════════════
-# 2. GAE Computation
-# ════════════════════════════════════════════════════════════════════════════════
-
-def test_gae():
-    print("\n--- GAE Computation ---")
-    # With V=0, GAE reduces to discounted returns
-    rewards = [1.0, 1.0, 1.0]
-    values = [0.0, 0.0, 0.0]
-    gae = compute_gae(rewards, values, gamma=0.99, lam=1.0)
-    returns = compute_returns(rewards, gamma=0.99)
-    record("GAE with V=0, lam=1 equals returns",
-           PASS if all(abs(g - r) < 0.01 for g, r in zip(gae, returns)) else FAIL,
-           f"gae={[round(g,2) for g in gae]} returns={[round(r,2) for r in returns]}")
-
-    # With lam=0, GAE reduces to one-step TD
-    gae_td = compute_gae(rewards, values, gamma=0.99, lam=0.0)
-    # δ_0 = r_0 + γ*V(s_1) - V(s_0) = 1.0 + 0.99*0 - 0 = 1.0
-    record("GAE with lam=0 is one-step TD",
-           PASS if abs(gae_td[0] - 1.0) < 0.001 else FAIL,
-           f"got {gae_td[0]:.4f} expected 1.0")
-
-    # With non-zero values, advantage is return minus value
-    values2 = [2.0, 1.0, 0.5]
-    gae2 = compute_gae([1.0, 1.0, 1.0], values2, gamma=0.99, lam=1.0)
-    record("GAE with values produces non-trivial advantages",
-           PASS if len(gae2) == 3 and any(abs(g) > 0.1 for g in gae2) else FAIL,
-           f"gae={[round(g,3) for g in gae2]}")
-
-
-# ════════════════════════════════════════════════════════════════════════════════
-# 3. REINFORCE Loss
-# ═══════════════════════════════════════════════════════════���════════════════════
-
-def test_reinforce():
-    print("\n--- REINFORCE Loss ---")
-    # Known values: log_probs=[-1, -2, -3], returns=[10, 5, 1]
-    # loss = -(1/3) * ((-1)*10 + (-2)*5 + (-3)*1) = -(1/3)*(-10-10-3) = -(1/3)*(-23) = 7.667
-    log_probs = [-1.0, -2.0, -3.0]
-    returns = [10.0, 5.0, 1.0]
-    loss = reinforce_loss(log_probs, returns)
-    expected = -((-1.0)*10 + (-2.0)*5 + (-3.0)*1) / 3
-    record("REINFORCE loss matches manual calculation",
-           PASS if abs(loss - expected) < 0.001 else FAIL,
-           f"got {loss:.4f} expected {expected:.4f}")
-
-    # Higher returns should produce higher loss (more gradient push)
-    returns_high = [20.0, 10.0, 2.0]
-    loss_high = reinforce_loss(log_probs, returns_high)
-    record("Higher returns → higher loss magnitude",
-           PASS if abs(loss_high) > abs(loss) else FAIL,
-           f"low={abs(loss):.4f} high={abs(loss_high):.4f}")
-
-    # Empty episode
-    loss_empty = reinforce_loss([], [])
-    record("REINFORCE handles empty episode",
-           PASS if loss_empty == 0.0 else FAIL,
-           f"got {loss_empty}")
-
-
-# ════════════════════════════════════════════════════════════════════════════════
-# 4. REINFORCE + Baseline Loss
-# ════════════════════════════════════════════════════════════════════════════════
-
-def test_reinforce_baseline():
-    print("\n--- REINFORCE + Baseline Loss ---")
-    log_probs = [-1.0, -2.0, -3.0]
-    returns = [10.0, 5.0, 1.0]
-
-    # With baselines=None, uses mean(returns)=5.33 as baseline
-    loss_b = reinforce_baseline_loss(log_probs, returns, baselines=None, normalize_advantage=False)
-    # advantages = [10-5.33, 5-5.33, 1-5.33] = [4.67, -0.33, -4.33]
-    # loss = -(1/3) * ((-1)*4.67 + (-2)*(-0.33) + (-3)*(-4.33))
-    #      = -(1/3) * (-4.67 + 0.67 + 13.0)
-    #      = -(1/3) * 9.0 = -3.0
-    mean_r = sum(returns) / len(returns)
-    advantages = [g - mean_r for g in returns]
-    expected = -sum(lp * adv for lp, adv in zip(log_probs, advantages)) / 3
-    record("REINFORCE+baseline matches manual calc",
-           PASS if abs(loss_b - expected) < 0.01 else FAIL,
-           f"got {loss_b:.4f} expected {expected:.4f}")
-
-    # With normalize_advantage=True, advantages are standardized
-    loss_norm = reinforce_baseline_loss(log_probs, returns, baselines=None, normalize_advantage=True)
-    record("Normalized advantage produces valid loss",
-           PASS if not math.isnan(loss_norm) and not math.isinf(loss_norm) else FAIL,
-           f"loss={loss_norm:.4f}")
-
-    # Baseline should reduce loss magnitude vs vanilla REINFORCE
-    loss_vanilla = reinforce_loss(log_probs, returns)
-    record("Baseline typically reduces loss magnitude",
-           PASS if abs(loss_norm) < abs(loss_vanilla) or True else FAIL,
-           f"vanilla={abs(loss_vanilla):.4f} baseline={abs(loss_norm):.4f} (varies)")
-
-    # Custom baselines
-    baselines = [9.0, 4.0, 0.5]
-    loss_custom = reinforce_baseline_loss(log_probs, returns, baselines=baselines, normalize_advantage=False)
-    advantages_custom = [g - b for g, b in zip(returns, baselines)]
-    expected_custom = -sum(lp * adv for lp, adv in zip(log_probs, advantages_custom)) / 3
-    record("Custom baselines work correctly",
-           PASS if abs(loss_custom - expected_custom) < 0.01 else FAIL,
-           f"got {loss_custom:.4f} expected {expected_custom:.4f}")
-
-
-# ════════════════════════════════════════════════════════════════════════════════
-# 5. GRPO Loss
-# ════════════════════════════════════════════════════════════════════════════════
-
-def test_grpo():
-    print("\n--- GRPO Loss ---")
-    # Group of 3 samples for one state
-    log_probs_groups = [[-1.0, -2.0, -1.5]]
-    rewards_groups = [[10.0, 5.0, 8.0]]
-
-    loss = grpo_loss(log_probs_groups, rewards_groups)
-    record("GRPO produces valid loss",
-           PASS if not math.isnan(loss) and not math.isinf(loss) else FAIL,
-           f"loss={loss:.4f}")
-
-    # The highest-reward sample should get positive advantage,
-    # lowest-reward should get negative advantage
-    mean_r = sum(rewards_groups[0]) / 3  # 7.67
-    std_r = math.sqrt(sum((r - mean_r)**2 for r in rewards_groups[0]) / 3)
-    advantages = [(r - mean_r) / (std_r + 1e-8) for r in rewards_groups[0]]
-    record("GRPO: highest reward gets positive advantage",
-           PASS if advantages[0] > 0 else FAIL,
-           f"adv={advantages[0]:.4f}")
-    record("GRPO: lowest reward gets negative advantage",
-           PASS if advantages[1] < 0 else FAIL,
-           f"adv={advantages[1]:.4f}")
-
-    # Multiple groups
-    log_probs_2 = [[-1.0, -2.0], [-1.5, -1.5]]
-    rewards_2 = [[10.0, 5.0], [3.0, 7.0]]
-    loss_2 = grpo_loss(log_probs_2, rewards_2)
-    record("GRPO handles multiple groups",
-           PASS if not math.isnan(loss_2) else FAIL,
-           f"loss={loss_2:.4f}")
-
-    # Empty groups
-    loss_empty = grpo_loss([], [])
-    record("GRPO handles empty input",
-           PASS if loss_empty == 0.0 else FAIL,
-           f"got {loss_empty}")
-
-    # Identical rewards → zero advantage → zero loss
-    loss_identical = grpo_loss([[-1.0, -2.0, -3.0]], [[5.0, 5.0, 5.0]])
-    record("GRPO: identical rewards → near-zero loss",
-           PASS if abs(loss_identical) < 1e-4 else FAIL,
-           f"loss={loss_identical:.6f}")
-
-
-# ════════════════════════════════════════════════════════════════════════════════
-# 6. RLOO Loss
-# ════════════════════════════════════════════════════════════════════════════════
-
-def test_rloo():
-    print("\n--- RLOO Loss ---")
-    # Group of 3 samples
-    log_probs_groups = [[-1.0, -2.0, -1.5]]
-    rewards_groups = [[10.0, 5.0, 8.0]]
-
-    loss = rloo_loss(log_probs_groups, rewards_groups)
-    record("RLOO produces valid loss",
-           PASS if not math.isnan(loss) and not math.isinf(loss) else FAIL,
-           f"loss={loss:.4f}")
-
-    # Leave-one-out baselines
-    # For r=10: baseline = (5+8)/2 = 6.5, advantage = 10-6.5 = 3.5
-    # For r=5: baseline = (10+8)/2 = 9.0, advantage = 5-9.0 = -4.0
-    # For r=8: baseline = (10+5)/2 = 7.5, advantage = 8-7.5 = 0.5
-    baselines = [6.5, 9.0, 7.5]
-    advantages = [10-6.5, 5-9.0, 8-7.5]
-    expected = -sum(lp * adv for lp, adv in zip(log_probs_groups[0], advantages)) / 3
-    record("RLOO matches manual calculation",
-           PASS if abs(loss - expected) < 0.01 else FAIL,
-           f"got {loss:.4f} expected {expected:.4f}")
-
-    # Single sample: falls back to REINFORCE
-    loss_single = rloo_loss([[-1.0]], [[5.0]])
-    expected_single = -(-1.0) * 5.0  # REINFORCE on one sample
-    record("RLOO K=1 falls back to REINFORCE",
-           PASS if abs(loss_single - expected_single) < 0.01 else FAIL,
-           f"got {loss_single:.4f} expected {expected_single:.4f}")
-
-    # K=2: simplest meaningful RLOO
-    loss_k2 = rloo_loss([[-1.0, -2.0]], [[10.0, 5.0]])
-    # baseline for r=10: 5.0, adv=5.0
-    # baseline for r=5: 10.0, adv=-5.0
-    # loss = -(1/2) * ((-1)*5 + (-2)*(-5)) = -(1/2)*(-5+10) = -2.5
-    record("RLOO K=2 produces valid loss",
-           PASS if not math.isnan(loss_k2) else FAIL,
-           f"loss={loss_k2:.4f}")
-
-
-# ════════════════════════════════════════════════════════════════════════════════
-# 7. Reward Normalization
-# ════════════════════════════════════════════════════════════════════════════════
-
-def test_reward_normalization():
-    print("\n--- Reward Normalization ---")
-    raw = [-0.5, -1.0, -0.3, -2.0, -0.8]
-    mean, var = compute_reward_stats(raw)
-    record("Reward stats computed",
-           PASS if abs(mean - (-0.92)) < 0.01 else FAIL,
-           f"mean={mean:.4f} var={var:.4f}")
-
-    normed = normalize_rewards(raw, mean, var)
-    record("Normalized rewards have near-zero mean",
-           PASS if abs(sum(normed)/len(normed)) < 0.01 else FAIL,
-           f"mean={sum(normed)/len(normed):.4f}")
-
-    norm_var = sum((n - sum(normed)/len(normed))**2 for n in normed) / len(normed)
-    record("Normalized rewards have near-unit variance",
-           PASS if abs(norm_var - 1.0) < 0.01 else FAIL,
-           f"var={norm_var:.4f}")
-
-    # Identity: normalizing with mean=0, var=1 should leave rewards unchanged
-    identity = normalize_rewards(raw, 0.0, 1.0)
-    record("Identity normalization (mean=0, var=1)",
-           PASS if all(abs(a - b) < 0.01 for a, b in zip(raw, identity)) else FAIL,
-           f"max_diff={max(abs(a-b) for a,b in zip(raw,identity)):.4f}")
-
-
-# ════════════════════════════════════════════════════════════════════════════════
-# 8. Loss Function Comparison
-# ════════════════════════════════════════════════════════════════════════════════
-
-def test_loss_comparison():
-    """Compare all 4 loss functions on the same episode data."""
-    print("\n--- Loss Function Comparison ---")
-    log_probs = [-2.0, -1.5, -3.0, -1.0, -2.5]
-    returns = [0.8, 0.3, 0.1, 0.5, 0.2]
-
-    l_reinforce = reinforce_loss(log_probs, returns)
-    l_baseline = reinforce_baseline_loss(log_probs, returns, normalize_advantage=True)
-
-    # GRPO: treat each step as its own "group" of size 1
-    # (Not how GRPO is normally used, but tests the pipeline)
-    lps_groups = [[lp] for lp in log_probs]
-    rs_groups = [[r] for r in returns]
-    l_grpo = grpo_loss(lps_groups, rs_groups)
-    l_rloo = rloo_loss(lps_groups, rs_groups)
-
-    record("All 4 losses produce valid values",
-           PASS if all(not math.isnan(l) and not math.isinf(l)
-                       for l in [l_reinforce, l_baseline, l_grpo, l_rloo]) else FAIL,
-           f"R={l_reinforce:.4f} RB={l_baseline:.4f} GRPO={l_grpo:.4f} RLOO={l_rloo:.4f}")
-
-    print(f"  [i] REINFORCE:           {l_reinforce:.6f}")
-    print(f"  [i] REINFORCE+baseline:  {l_baseline:.6f}")
-    print(f"  [i] GRPO (K=1):          {l_grpo:.6f}")
-    print(f"  [i] RLOO (K=1):          {l_rloo:.6f}")
-
-    # Now with proper K=4 groups
-    log_probs_4 = [[-1.0, -2.0, -1.5, -3.0]]
-    rewards_4 = [[0.8, 0.2, 0.5, 0.1]]
-    l_grpo_4 = grpo_loss(log_probs_4, rewards_4)
-    l_rloo_4 = rloo_loss(log_probs_4, rewards_4)
-    record("GRPO/RLOO with K=4 produce valid losses",
-           PASS if not math.isnan(l_grpo_4) and not math.isnan(l_rloo_4) else FAIL,
-           f"GRPO={l_grpo_4:.4f} RLOO={l_rloo_4:.4f}")
-    print(f"  [i] GRPO (K=4): {l_grpo_4:.6f}")
-    print(f"  [i] RLOO (K=4): {l_rloo_4:.6f}")
-
-
-# ════════════════════════════════════════════════════════════════════════════════
-# 9. Episode Collection (with MockPolicyModel)
-# ════════════════════════════════════════════════════════════════════════════════
-
-def test_episode_collection():
-    print("\n--- Episode Collection (MockPolicyModel) ---")
-    config = TrainingConfig(n_nodes=5, max_steps=30)
-    collector = EpisodeCollector(config)
-    model = MockPolicyModel(n_nodes=5, seed=42)
-
-    episode = collector.collect_episode(model, task_id="task-1", seed=42)
-
-    record("Episode has correct number of steps",
-           PASS if len(episode.steps) == 30 else FAIL,
-           f"steps={len(episode.steps)}")
-
-    record("All log probs are valid",
-           PASS if all(not math.isnan(s.log_prob) for s in episode.steps) else FAIL,
-           f"min_lp={min(s.log_prob for s in episode.steps):.4f}")
-
-    record("Rewards are finite",
-           PASS if all(math.isfinite(s.reward) for s in episode.steps) else FAIL,
-           f"min_r={min(s.reward for s in episode.steps):.4f}")
-
-    record("Normalized rewards in [0,1]",
-           PASS if all(0.0 <= s.reward_normalized <= 1.0 for s in episode.steps) else FAIL,
-           f"range=[{min(s.reward_normalized for s in episode.steps):.4f}, "
-           f"{max(s.reward_normalized for s in episode.steps):.4f}]")
-
-    record("Total reward is computed",
-           PASS if math.isfinite(episode.total_reward) else FAIL,
-           f"total={episode.total_reward:.4f}")
-
-    record("SLA violations tracked",
-           PASS if isinstance(episode.sla_violations, int) else FAIL,
-           f"violations={episode.sla_violations}")
-
-
-# ════════════════════════════════════════════════════════════════════════════════
-# 10. Full Training Step (per loss function)
-# ════════════════════════════════════════════════════════════════════════════════
-
-def test_training_steps():
-    """Run one training step with each loss function."""
-    print("\n--- Full Training Steps ---")
-    model = MockPolicyModel(n_nodes=5, seed=42)
-
-    for loss_name in [LOSS_REINFORCE, LOSS_REINFORCE_BASELINE, LOSS_GRPO, LOSS_RLOO]:
-        config = TrainingConfig(
-            n_nodes=5,
-            max_steps=30,
-            loss_fn=loss_name,
-            n_samples_per_state=2 if loss_name in (LOSS_GRPO, LOSS_RLOO) else 1,
-        )
-        trainer = SRETrainer(config)
-        metrics = trainer.train_step(model, task_id="task-1", seed=42)
-
-        record(f"{loss_name}: loss is valid",
-               PASS if math.isfinite(metrics["loss"]) else FAIL,
-               f"loss={metrics['loss']:.4f}")
-
-        record(f"{loss_name}: avg_reward is valid",
-               PASS if math.isfinite(metrics["avg_reward"]) else FAIL,
-               f"avg_reward={metrics['avg_reward']:.4f}")
-
-        record(f"{loss_name}: episode completed",
-               PASS if metrics["episode_length"] > 0 else FAIL,
-               f"length={metrics['episode_length']}")
-
-        # No NaN/inf in running stats
-        record(f"{loss_name}: running stats stable",
-               PASS if math.isfinite(metrics["reward_mean"]) and math.isfinite(metrics["reward_var"]) else FAIL,
-               f"mean={metrics['reward_mean']:.4f} var={metrics['reward_var']:.4f}")
-
-
-# ════════════════════════════════════════════════════════════════════════════════
-# 11. Multi-Episode Stability
-# ════════════════════════════════════════════════════════════════════════════════
-
-def test_multi_episode_stability():
-    """Run multiple episodes and check running stats remain stable."""
-    print("\n--- Multi-Episode Stability ---")
-    config = TrainingConfig(
-        n_nodes=5,
-        max_steps=30,
-        loss_fn=LOSS_REINFORCE_BASELINE,
-        tasks=["task-1", "task-2", "task-3"],
-    )
-    trainer = SRETrainer(config)
-    model = MockPolicyModel(n_nodes=5, seed=42)
-
-    all_losses = []
-    for i in range(5):
-        for task in config.tasks:
-            metrics = trainer.train_step(model, task_id=task, seed=42 + i)
-            all_losses.append(metrics["loss"])
-
-    # No NaN/inf across 15 episodes
-    record("15 episodes: all losses finite",
-           PASS if all(math.isfinite(l) for l in all_losses) else FAIL,
-           f"n_losses={len(all_losses)}")
-
-    # Losses should vary (different tasks + domain randomization)
-    unique_losses = len(set(round(l, 4) for l in all_losses))
-    record("Losses vary across episodes",
-           PASS if unique_losses > 3 else FAIL,
-           f"unique={unique_losses}/{len(all_losses)}")
-
-    # Running stats should be non-degenerate
-    last_metrics = metrics
-    record("Running reward mean is non-zero",
-           PASS if abs(last_metrics["reward_mean"]) > 0.001 else FAIL,
-           f"mean={last_metrics['reward_mean']:.6f}")
-
-
-# ════════════════════════════════════════════════════════════════════════════════
-# 12. SRE-Specific Edge Cases
-# ════════════════════════════════════════════════════════════════════════════════
-
-def test_sre_edge_cases():
-    """Test edge cases specific to the SRE domain."""
-    print("\n--- SRE Edge Cases ---")
-
-    # Very negative rewards (system crashing)
-    log_probs = [-2.0] * 10
-    returns_crash = [-100.0] * 10
-    loss_crash = reinforce_baseline_loss(log_probs, returns_crash, normalize_advantage=True)
-    record("Very negative rewards: loss is finite",
-           PASS if math.isfinite(loss_crash) else FAIL,
-           f"loss={loss_crash:.4f}")
-
-    # All-zero returns (perfect episode)
-    returns_perfect = [0.0] * 10
-    loss_perfect = reinforce_baseline_loss(log_probs, returns_perfect, normalize_advantage=False)
-    record("Zero returns: loss is zero (no gradient)",
-           PASS if abs(loss_perfect) < 0.001 else FAIL,
-           f"loss={loss_perfect:.4f}")
-
-    # Highly variable rewards within episode (surge task)
-    returns_surge = [0.5, 0.5, -10.0, -10.0, 0.5, 0.5, -10.0, 0.5, 0.5, 0.5]
-    loss_surge = reinforce_baseline_loss(log_probs, returns_surge, normalize_advantage=True)
-    record("High-variance rewards: loss is finite with normalization",
-           PASS if math.isfinite(loss_surge) else FAIL,
-           f"loss={loss_surge:.4f}")
-
-    # GRPO with very different rewards in group
-    lps = [[-1.0, -2.0, -1.5, -3.0]]
-    rs_extreme = [[0.9, 0.8, 0.85, 0.05]]  # One bad sample
-    loss_extreme = grpo_loss(lps, rs_extreme)
-    record("GRPO handles outlier in group",
-           PASS if math.isfinite(loss_extreme) else FAIL,
-           f"loss={loss_extreme:.4f}")
-
-
-# ════════════════════════════════════════════════════════════════════════════════
-
-def main():
-    print("=" * 60)
-    print("AntiAtropos Training Validation")
-    print("=" * 60)
-
-    test_returns()
-    test_gae()
-    test_reinforce()
-    test_reinforce_baseline()
-    test_grpo()
-    test_rloo()
-    test_reward_normalization()
-    test_loss_comparison()
-    test_episode_collection()
-    test_training_steps()
-    test_multi_episode_stability()
-    test_sre_edge_cases()
-
-    passed = sum(1 for _, s, _ in results if s == PASS)
-    failed = sum(1 for _, s, _ in results if s == FAIL)
-    total = len(results)
-
-    print("\n" + "=" * 60)
-    print(f"RESULTS: {passed}/{total} passed, {failed} failed")
-    print("=" * 60)
-
-    if failed > 0:
-        print("\nFailed tests:")
-        for name, status, detail in results:
-            if status == FAIL:
-                print(f"  X {name}: {detail}")
-
-    return 0 if failed == 0 else 1
-
-
-if __name__ == "__main__":
-    sys.exit(main())
+#!/usr/bin/env python3
+"""
+AntiAtropos Training Validation — Local Test.
+
+Validates the training pipeline (loss functions, episode collection, trainer)
+using a MockPolicyModel (no GPU needed). Run before going to Colab.
+
+Run from project root:
+    python train_test.py
+"""
+
+import sys
+import os
+import math
+import random
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+from training.losses import (
+    compute_returns, compute_gae,
+    reinforce_loss, reinforce_baseline_loss,
+    grpo_loss, rloo_loss,
+    normalize_rewards, compute_reward_stats,
+)
+from training.trainer import (
+    SRETrainer, TrainingConfig, EpisodeCollector,
+    MockPolicyModel,
+    LOSS_REINFORCE, LOSS_REINFORCE_BASELINE, LOSS_GRPO, LOSS_RLOO,
+)
+
+PASS = "PASS"
+FAIL = "FAIL"
+results: list[tuple[str, str, str]] = []
+
+
+def record(name: str, status: str, detail: str = "") -> None:
+    results.append((name, status, detail))
+    icon = "+" if status == PASS else "X"
+    msg = f"  [{icon}] {name}"
+    if detail:
+        msg += f" -- {detail}"
+    print(msg)
+
+
+# ════════════════════════════════════════════════════════════════════════════════
+# 1. Return Computation
+# ════════════════════════════════════════════════════════════════════════════════
+
+def test_returns():
+    print("\n--- Return Computation ---")
+    # Simple case: [1, 1, 1] with gamma=0.99
+    returns = compute_returns([1.0, 1.0, 1.0], gamma=0.99)
+    # G_2 = 1.0, G_1 = 1 + 0.99*1 = 1.99, G_0 = 1 + 0.99*1.99 = 2.9701
+    record("compute_returns[0]",
+           PASS if abs(returns[0] - 2.9701) < 0.001 else FAIL,
+           f"got {returns[0]:.4f} expected 2.9701")
+    record("compute_returns[2]",
+           PASS if abs(returns[2] - 1.0) < 0.001 else FAIL,
+           f"got {returns[2]:.4f} expected 1.0")
+
+    # Empty rewards
+    returns_empty = compute_returns([])
+    record("compute_returns handles empty",
+           PASS if returns_empty == [] else FAIL,
+           f"got {returns_empty}")
+
+    # Single reward
+    returns_single = compute_returns([5.0])
+    record("compute_returns single reward",
+           PASS if abs(returns_single[0] - 5.0) < 0.001 else FAIL,
+           f"got {returns_single[0]:.4f}")
+
+    # Discount factor = 0 → only immediate reward matters
+    returns_0 = compute_returns([1.0, 2.0, 3.0], gamma=0.0)
+    record("gamma=0: returns = rewards",
+           PASS if returns_0 == [1.0, 2.0, 3.0] else FAIL,
+           f"got {returns_0}")
+
+
+# ════════════════════════════════════════════════════════════════════════════════
+# 2. GAE Computation
+# ════════════════════════════════════════════════════════════════════════════════
+
+def test_gae():
+    print("\n--- GAE Computation ---")
+    # With V=0, GAE reduces to discounted returns
+    rewards = [1.0, 1.0, 1.0]
+    values = [0.0, 0.0, 0.0]
+    gae = compute_gae(rewards, values, gamma=0.99, lam=1.0)
+    returns = compute_returns(rewards, gamma=0.99)
+    record("GAE with V=0, lam=1 equals returns",
+           PASS if all(abs(g - r) < 0.01 for g, r in zip(gae, returns)) else FAIL,
+           f"gae={[round(g,2) for g in gae]} returns={[round(r,2) for r in returns]}")
+
+    # With lam=0, GAE reduces to one-step TD
+    gae_td = compute_gae(rewards, values, gamma=0.99, lam=0.0)
+    # δ_0 = r_0 + γ*V(s_1) - V(s_0) = 1.0 + 0.99*0 - 0 = 1.0
+    record("GAE with lam=0 is one-step TD",
+           PASS if abs(gae_td[0] - 1.0) < 0.001 else FAIL,
+           f"got {gae_td[0]:.4f} expected 1.0")
+
+    # With non-zero values, advantage is return minus value
+    values2 = [2.0, 1.0, 0.5]
+    gae2 = compute_gae([1.0, 1.0, 1.0], values2, gamma=0.99, lam=1.0)
+    record("GAE with values produces non-trivial advantages",
+           PASS if len(gae2) == 3 and any(abs(g) > 0.1 for g in gae2) else FAIL,
+           f"gae={[round(g,3) for g in gae2]}")
+
+
+# ════════════════════════════════════════════════════════════════════════════════
+# 3. REINFORCE Loss
+# ════════════════════════════════════════════════════════════════════════════════
+
+def test_reinforce():
+    print("\n--- REINFORCE Loss ---")
+    # Known values: log_probs=[-1, -2, -3], returns=[10, 5, 1]
+    # loss = -(1/3) * ((-1)*10 + (-2)*5 + (-3)*1) = -(1/3)*(-10-10-3) = -(1/3)*(-23) = 7.667
+    log_probs = [-1.0, -2.0, -3.0]
+    returns = [10.0, 5.0, 1.0]
+    loss = reinforce_loss(log_probs, returns)
+    expected = -((-1.0)*10 + (-2.0)*5 + (-3.0)*1) / 3
+    record("REINFORCE loss matches manual calculation",
+           PASS if abs(loss - expected) < 0.001 else FAIL,
+           f"got {loss:.4f} expected {expected:.4f}")
+
+    # Higher returns should produce higher loss (more gradient push)
+    returns_high = [20.0, 10.0, 2.0]
+    loss_high = reinforce_loss(log_probs, returns_high)
+    record("Higher returns → higher loss magnitude",
+           PASS if abs(loss_high) > abs(loss) else FAIL,
+           f"low={abs(loss):.4f} high={abs(loss_high):.4f}")
+
+    # Empty episode
+    loss_empty = reinforce_loss([], [])
+    record("REINFORCE handles empty episode",
+           PASS if loss_empty == 0.0 else FAIL,
+           f"got {loss_empty}")
+
+
+# ════════════════════════════════════════════════════════════════════════════════
+# 4. REINFORCE + Baseline Loss
+# ════════════════════════════════════════════════════════════════════════════════
+
+def test_reinforce_baseline():
+    print("\n--- REINFORCE + Baseline Loss ---")
+    log_probs = [-1.0, -2.0, -3.0]
+    returns = [10.0, 5.0, 1.0]
+
+    # With baselines=None, uses mean(returns)=5.33 as baseline
+    loss_b = reinforce_baseline_loss(log_probs, returns, baselines=None, normalize_advantage=False)
+    # advantages = [10-5.33, 5-5.33, 1-5.33] = [4.67, -0.33, -4.33]
+    # loss = -(1/3) * ((-1)*4.67 + (-2)*(-0.33) + (-3)*(-4.33))
+    #      = -(1/3) * (-4.67 + 0.67 + 13.0)
+    #      = -(1/3) * 9.0 = -3.0
+    mean_r = sum(returns) / len(returns)
+    advantages = [g - mean_r for g in returns]
+    expected = -sum(lp * adv for lp, adv in zip(log_probs, advantages)) / 3
+    record("REINFORCE+baseline matches manual calc",
+           PASS if abs(loss_b - expected) < 0.01 else FAIL,
+           f"got {loss_b:.4f} expected {expected:.4f}")
+
+    # With normalize_advantage=True, advantages are standardized
+    loss_norm = reinforce_baseline_loss(log_probs, returns, baselines=None, normalize_advantage=True)
+    record("Normalized advantage produces valid loss",
+           PASS if not math.isnan(loss_norm) and not math.isinf(loss_norm) else FAIL,
+           f"loss={loss_norm:.4f}")
+
+    # Baseline should reduce loss magnitude vs vanilla REINFORCE
+    loss_vanilla = reinforce_loss(log_probs, returns)
+    record("Baseline typically reduces loss magnitude",
+           PASS if abs(loss_norm) < abs(loss_vanilla) or True else FAIL,
+           f"vanilla={abs(loss_vanilla):.4f} baseline={abs(loss_norm):.4f} (varies)")
+
+    # Custom baselines
+    baselines = [9.0, 4.0, 0.5]
+    loss_custom = reinforce_baseline_loss(log_probs, returns, baselines=baselines, normalize_advantage=False)
+    advantages_custom = [g - b for g, b in zip(returns, baselines)]
+    expected_custom = -sum(lp * adv for lp, adv in zip(log_probs, advantages_custom)) / 3
+    record("Custom baselines work correctly",
+           PASS if abs(loss_custom - expected_custom) < 0.01 else FAIL,
+           f"got {loss_custom:.4f} expected {expected_custom:.4f}")
+
+
+# ════════════════════════════════════════════════════════════════════════════════
+# 5. GRPO Loss
+# ════════════════════════════════════��═══════════════════════════════════════════
+
+def test_grpo():
+    print("\n--- GRPO Loss ---")
+    # Group of 3 samples for one state
+    log_probs_groups = [[-1.0, -2.0, -1.5]]
+    rewards_groups = [[10.0, 5.0, 8.0]]
+
+    loss = grpo_loss(log_probs_groups, rewards_groups)
+    record("GRPO produces valid loss",
+           PASS if not math.isnan(loss) and not math.isinf(loss) else FAIL,
+           f"loss={loss:.4f}")
+
+    # The highest-reward sample should get positive advantage,
+    # lowest-reward should get negative advantage
+    mean_r = sum(rewards_groups[0]) / 3  # 7.67
+    std_r = math.sqrt(sum((r - mean_r)**2 for r in rewards_groups[0]) / 3)
+    advantages = [(r - mean_r) / (std_r + 1e-8) for r in rewards_groups[0]]
+    record("GRPO: highest reward gets positive advantage",
+           PASS if advantages[0] > 0 else FAIL,
+           f"adv={advantages[0]:.4f}")
+    record("GRPO: lowest reward gets negative advantage",
+           PASS if advantages[1] < 0 else FAIL,
+           f"adv={advantages[1]:.4f}")
+
+    # Multiple groups
+    log_probs_2 = [[-1.0, -2.0], [-1.5, -1.5]]
+    rewards_2 = [[10.0, 5.0], [3.0, 7.0]]
+    loss_2 = grpo_loss(log_probs_2, rewards_2)
+    record("GRPO handles multiple groups",
+           PASS if not math.isnan(loss_2) else FAIL,
+           f"loss={loss_2:.4f}")
+
+    # Empty groups
+    loss_empty = grpo_loss([], [])
+    record("GRPO handles empty input",
+           PASS if loss_empty == 0.0 else FAIL,
+           f"got {loss_empty}")
+
+    # Identical rewards → zero advantage → zero loss
+    loss_identical = grpo_loss([[-1.0, -2.0, -3.0]], [[5.0, 5.0, 5.0]])
+    record("GRPO: identical rewards → near-zero loss",
+           PASS if abs(loss_identical) < 1e-4 else FAIL,
+           f"loss={loss_identical:.6f}")
+
+
+# ════════════════════════════════════════════════════════════════════════════════
+# 6. RLOO Loss
+# ════════════════════════════════════════════════════════════════════════════════
+
+def test_rloo():
+    print("\n--- RLOO Loss ---")
+    # Group of 3 samples
+    log_probs_groups = [[-1.0, -2.0, -1.5]]
+    rewards_groups = [[10.0, 5.0, 8.0]]
+
+    loss = rloo_loss(log_probs_groups, rewards_groups)
+    record("RLOO produces valid loss",
+           PASS if not math.isnan(loss) and not math.isinf(loss) else FAIL,
+           f"loss={loss:.4f}")
+
+    # Leave-one-out baselines
+    # For r=10: baseline = (5+8)/2 = 6.5, advantage = 10-6.5 = 3.5
+    # For r=5: baseline = (10+8)/2 = 9.0, advantage = 5-9.0 = -4.0
+    # For r=8: baseline = (10+5)/2 = 7.5, advantage = 8-7.5 = 0.5
+    baselines = [6.5, 9.0, 7.5]
+    advantages = [10-6.5, 5-9.0, 8-7.5]
+    expected = -sum(lp * adv for lp, adv in zip(log_probs_groups[0], advantages)) / 3
+    record("RLOO matches manual calculation",
+           PASS if abs(loss - expected) < 0.01 else FAIL,
+           f"got {loss:.4f} expected {expected:.4f}")
+
+    # Single sample: falls back to REINFORCE
+    loss_single = rloo_loss([[-1.0]], [[5.0]])
+    expected_single = -(-1.0) * 5.0  # REINFORCE on one sample
+    record("RLOO K=1 falls back to REINFORCE",
+           PASS if abs(loss_single - expected_single) < 0.01 else FAIL,
+           f"got {loss_single:.4f} expected {expected_single:.4f}")
+
+    # K=2: simplest meaningful RLOO
+    loss_k2 = rloo_loss([[-1.0, -2.0]], [[10.0, 5.0]])
+    # baseline for r=10: 5.0, adv=5.0
+    # baseline for r=5: 10.0, adv=-5.0
+    # loss = -(1/2) * ((-1)*5 + (-2)*(-5)) = -(1/2)*(-5+10) = -2.5
+    record("RLOO K=2 produces valid loss",
+           PASS if not math.isnan(loss_k2) else FAIL,
+           f"loss={loss_k2:.4f}")
+
+
+# ════════════════════════════════════════════════════════════════════════════════
+# 7. Reward Normalization
+# ════════════════════════════════════════════════════════════════════════════════
+
+def test_reward_normalization():
+    print("\n--- Reward Normalization ---")
+    raw = [-0.5, -1.0, -0.3, -2.0, -0.8]
+    mean, var = compute_reward_stats(raw)
+    record("Reward stats computed",
+           PASS if abs(mean - (-0.92)) < 0.01 else FAIL,
+           f"mean={mean:.4f} var={var:.4f}")
+
+    normed = normalize_rewards(raw, mean, var)
+    record("Normalized rewards have near-zero mean",
+           PASS if abs(sum(normed)/len(normed)) < 0.01 else FAIL,
+           f"mean={sum(normed)/len(normed):.4f}")
+
+    norm_var = sum((n - sum(normed)/len(normed))**2 for n in normed) / len(normed)
+    record("Normalized rewards have near-unit variance",
+           PASS if abs(norm_var - 1.0) < 0.01 else FAIL,
+           f"var={norm_var:.4f}")
+
+    # Identity: normalizing with mean=0, var=1 should leave rewards unchanged
+    identity = normalize_rewards(raw, 0.0, 1.0)
+    record("Identity normalization (mean=0, var=1)",
+           PASS if all(abs(a - b) < 0.01 for a, b in zip(raw, identity)) else FAIL,
+           f"max_diff={max(abs(a-b) for a,b in zip(raw,identity)):.4f}")
+
+
+# ════════════════════════════════════════════════════════════════════════════════
+# 8. Loss Function Comparison
+# ════════════════════════════════════════════════════════════════════════════════
+
+def test_loss_comparison():
+    """Compare all 4 loss functions on the same episode data."""
+    print("\n--- Loss Function Comparison ---")
+    log_probs = [-2.0, -1.5, -3.0, -1.0, -2.5]
+    returns = [0.8, 0.3, 0.1, 0.5, 0.2]
+
+    l_reinforce = reinforce_loss(log_probs, returns)
+    l_baseline = reinforce_baseline_loss(log_probs, returns, normalize_advantage=True)
+
+    # GRPO: treat each step as its own "group" of size 1
+    # (Not how GRPO is normally used, but tests the pipeline)
+    lps_groups = [[lp] for lp in log_probs]
+    rs_groups = [[r] for r in returns]
+    l_grpo = grpo_loss(lps_groups, rs_groups)
+    l_rloo = rloo_loss(lps_groups, rs_groups)
+
+    record("All 4 losses produce valid values",
+           PASS if all(not math.isnan(l) and not math.isinf(l)
+                       for l in [l_reinforce, l_baseline, l_grpo, l_rloo]) else FAIL,
+           f"R={l_reinforce:.4f} RB={l_baseline:.4f} GRPO={l_grpo:.4f} RLOO={l_rloo:.4f}")
+
+    print(f"  [i] REINFORCE:           {l_reinforce:.6f}")
+    print(f"  [i] REINFORCE+baseline:  {l_baseline:.6f}")
+    print(f"  [i] GRPO (K=1):          {l_grpo:.6f}")
+    print(f"  [i] RLOO (K=1):          {l_rloo:.6f}")
+
+    # Now with proper K=4 groups
+    log_probs_4 = [[-1.0, -2.0, -1.5, -3.0]]
+    rewards_4 = [[0.8, 0.2, 0.5, 0.1]]
+    l_grpo_4 = grpo_loss(log_probs_4, rewards_4)
+    l_rloo_4 = rloo_loss(log_probs_4, rewards_4)
+    record("GRPO/RLOO with K=4 produce valid losses",
+           PASS if not math.isnan(l_grpo_4) and not math.isnan(l_rloo_4) else FAIL,
+           f"GRPO={l_grpo_4:.4f} RLOO={l_rloo_4:.4f}")
+    print(f"  [i] GRPO (K=4): {l_grpo_4:.6f}")
+    print(f"  [i] RLOO (K=4): {l_rloo_4:.6f}")
+
+
+# ════════════════════════════════════════════════════════════════════════════════
+# 9. Episode Collection (with MockPolicyModel)
+# ════════════════════════════════════════════════════════════════════════════════
+
+def test_episode_collection():
+    print("\n--- Episode Collection (MockPolicyModel) ---")
+    config = TrainingConfig(n_nodes=5, max_steps=30)
+    collector = EpisodeCollector(config)
+    model = MockPolicyModel(n_nodes=5, seed=42)
+
+    episode = collector.collect_episode(model, task_id="task-1", seed=42)
+
+    record("Episode has correct number of steps",
+           PASS if len(episode.steps) == 30 else FAIL,
+           f"steps={len(episode.steps)}")
+
+    record("All log probs are valid",
+           PASS if all(not math.isnan(s.log_prob) for s in episode.steps) else FAIL,
+           f"min_lp={min(s.log_prob for s in episode.steps):.4f}")
+
+    record("Rewards are finite",
+           PASS if all(math.isfinite(s.reward) for s in episode.steps) else FAIL,
+           f"min_r={min(s.reward for s in episode.steps):.4f}")
+
+    record("Normalized rewards in [0,1]",
+           PASS if all(0.0 <= s.reward_normalized <= 1.0 for s in episode.steps) else FAIL,
+           f"range=[{min(s.reward_normalized for s in episode.steps):.4f}, "
+           f"{max(s.reward_normalized for s in episode.steps):.4f}]")
+
+    record("Total reward is computed",
+           PASS if math.isfinite(episode.total_reward) else FAIL,
+           f"total={episode.total_reward:.4f}")
+
+    record("SLA violations tracked",
+           PASS if isinstance(episode.sla_violations, int) else FAIL,
+           f"violations={episode.sla_violations}")
+
+
+# ════════════════════════════════════════════════════════════════════════════════
+# 10. Full Training Step (per loss function)
+# ═══���════════════════════════════════════════════════════════════════════════════
+
+def test_training_steps():
+    """Run one training step with each loss function."""
+    print("\n--- Full Training Steps ---")
+    model = MockPolicyModel(n_nodes=5, seed=42)
+
+    for loss_name in [LOSS_REINFORCE, LOSS_REINFORCE_BASELINE, LOSS_GRPO, LOSS_RLOO]:
+        config = TrainingConfig(
+            n_nodes=5,
+            max_steps=30,
+            loss_fn=loss_name,
+            n_samples_per_state=2 if loss_name in (LOSS_GRPO, LOSS_RLOO) else 1,
+        )
+        trainer = SRETrainer(config)
+        metrics = trainer.train_step(model, task_id="task-1", seed=42)
+
+        record(f"{loss_name}: loss is valid",
+               PASS if math.isfinite(metrics["loss"]) else FAIL,
+               f"loss={metrics['loss']:.4f}")
+
+        record(f"{loss_name}: avg_reward is valid",
+               PASS if math.isfinite(metrics["avg_reward"]) else FAIL,
+               f"avg_reward={metrics['avg_reward']:.4f}")
+
+        record(f"{loss_name}: episode completed",
+               PASS if metrics["episode_length"] > 0 else FAIL,
+               f"length={metrics['episode_length']}")
+
+        # No NaN/inf in running stats
+        record(f"{loss_name}: running stats stable",
+               PASS if math.isfinite(metrics["reward_mean"]) and math.isfinite(metrics["reward_var"]) else FAIL,
+               f"mean={metrics['reward_mean']:.4f} var={metrics['reward_var']:.4f}")
+
+
+# ════════════════════════════════════════════════════════════════════════════════
+# 11. Multi-Episode Stability
+# ════════════════════════════════════════════════════════════════════════════════
+
+def test_multi_episode_stability():
+    """Run multiple episodes and check running stats remain stable."""
+    print("\n--- Multi-Episode Stability ---")
+    config = TrainingConfig(
+        n_nodes=5,
+        max_steps=30,
+        loss_fn=LOSS_REINFORCE_BASELINE,
+        tasks=["task-1", "task-2", "task-3"],
+    )
+    trainer = SRETrainer(config)
+    model = MockPolicyModel(n_nodes=5, seed=42)
+
+    all_losses = []
+    for i in range(5):
+        for task in config.tasks:
+            metrics = trainer.train_step(model, task_id=task, seed=42 + i)
+            all_losses.append(metrics["loss"])
+
+    # No NaN/inf across 15 episodes
+    record("15 episodes: all losses finite",
+           PASS if all(math.isfinite(l) for l in all_losses) else FAIL,
+           f"n_losses={len(all_losses)}")
+
+    # Losses should vary (different tasks + domain randomization)
+    unique_losses = len(set(round(l, 4) for l in all_losses))
+    record("Losses vary across episodes",
+           PASS if unique_losses > 3 else FAIL,
+           f"unique={unique_losses}/{len(all_losses)}")
+
+    # Running stats should be non-degenerate
+    last_metrics = metrics
+    record("Running reward mean is non-zero",
+           PASS if abs(last_metrics["reward_mean"]) > 0.001 else FAIL,
+           f"mean={last_metrics['reward_mean']:.6f}")
+
+
+# ════════════════════════════════════════════════════════════════════════════════
+# 12. SRE-Specific Edge Cases
+# ════════════════════════════════════════════════════════════════════════════════
+
+def test_sre_edge_cases():
+    """Test edge cases specific to the SRE domain."""
+    print("\n--- SRE Edge Cases ---")
+
+    # Very negative rewards (system crashing)
+    log_probs = [-2.0] * 10
+    returns_crash = [-100.0] * 10
+    loss_crash = reinforce_baseline_loss(log_probs, returns_crash, normalize_advantage=True)
+    record("Very negative rewards: loss is finite",
+           PASS if math.isfinite(loss_crash) else FAIL,
+           f"loss={loss_crash:.4f}")
+
+    # All-zero returns (perfect episode)
+    returns_perfect = [0.0] * 10
+    loss_perfect = reinforce_baseline_loss(log_probs, returns_perfect, normalize_advantage=False)
+    record("Zero returns: loss is zero (no gradient)",
+           PASS if abs(loss_perfect) < 0.001 else FAIL,
+           f"loss={loss_perfect:.4f}")
+
+    # Highly variable rewards within episode (surge task)
+    returns_surge = [0.5, 0.5, -10.0, -10.0, 0.5, 0.5, -10.0, 0.5, 0.5, 0.5]
+    loss_surge = reinforce_baseline_loss(log_probs, returns_surge, normalize_advantage=True)
+    record("High-variance rewards: loss is finite with normalization",
+           PASS if math.isfinite(loss_surge) else FAIL,
+           f"loss={loss_surge:.4f}")
+
+    # GRPO with very different rewards in group
+    lps = [[-1.0, -2.0, -1.5, -3.0]]
+    rs_extreme = [[0.9, 0.8, 0.85, 0.05]]  # One bad sample
+    loss_extreme = grpo_loss(lps, rs_extreme)
+    record("GRPO handles outlier in group",
+           PASS if math.isfinite(loss_extreme) else FAIL,
+           f"loss={loss_extreme:.4f}")
+
+
+# ════════════════════════════════════════════════════════════════════════════════
+
+def main():
+    print("=" * 60)
+    print("AntiAtropos Training Validation")
+    print("=" * 60)
+
+    test_returns()
+    test_gae()
+    test_reinforce()
+    test_reinforce_baseline()
+    test_grpo()
+    test_rloo()
+    test_reward_normalization()
+    test_loss_comparison()
+    test_episode_collection()
+    test_training_steps()
+    test_multi_episode_stability()
+    test_sre_edge_cases()
+
+    passed = sum(1 for _, s, _ in results if s == PASS)
+    failed = sum(1 for _, s, _ in results if s == FAIL)
+    total = len(results)
+
+    print("\n" + "=" * 60)
+    print(f"RESULTS: {passed}/{total} passed, {failed} failed")
+    print("=" * 60)
+
+    if failed > 0:
+        print("\nFailed tests:")
+        for name, status, detail in results:
+            if status == FAIL:
+                print(f"  X {name}: {detail}")
+
+    return 0 if failed == 0 else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

training/__init__.py

@@ -1,27 +1,27 @@
-"""AntiAtropos Training Module.
-
-RL loss functions and training loop for LLM-based SRE agents.
-Works with the simulator in pure-Python mode (no AWS/GPU needed for validation).
-"""
-
-from .losses import (
-    compute_returns,
-    compute_gae,
-    reinforce_loss,
-    reinforce_baseline_loss,
-    grpo_loss,
-    rloo_loss,
-)
-from .trainer import SRETrainer, TrainingConfig, EpisodeCollector
-
-__all__ = [
-    "compute_returns",
-    "compute_gae",
-    "reinforce_loss",
-    "reinforce_baseline_loss",
-    "grpo_loss",
-    "rloo_loss",
-    "SRETrainer",
-    "TrainingConfig",
-    "EpisodeCollector",
-]
+"""AntiAtropos Training Module.
+
+RL loss functions and training loop for LLM-based SRE agents.
+Works with the simulator in pure-Python mode (no AWS/GPU needed for validation).
+"""
+
+from .losses import (
+    compute_returns,
+    compute_gae,
+    reinforce_loss,
+    reinforce_baseline_loss,
+    grpo_loss,
+    rloo_loss,
+)
+from .trainer import SRETrainer, TrainingConfig, EpisodeCollector
+
+__all__ = [
+    "compute_returns",
+    "compute_gae",
+    "reinforce_loss",
+    "reinforce_baseline_loss",
+    "grpo_loss",
+    "rloo_loss",
+    "SRETrainer",
+    "TrainingConfig",
+    "EpisodeCollector",
+]

training/losses.py

@@ -1,390 +1,390 @@
-"""
-AntiAtropos RL Loss Functions.
-
-Pure-Python implementations of policy gradient loss functions for LLM-based
-SRE agents. These are mathematically identical to their PyTorch counterparts
-and can be validated locally without GPU.
-
-When porting to PyTorch (Colab), simply replace:
-    - float ops with torch.tensor ops
-    - sum() with torch.sum()
-    - log() with torch.log()
-    - The gradient flows through log_prob automatically
-
-Loss function selection guide for the SRE domain:
-─────────────────────────────────────────────────
-┌───────────────────────┬────────────┬──────────────┬─────────────────────────┐
-│ Method                │ Variance   │ Needs Value  │ Best for                │
-├───────────────────────┼────────────┼──────────────┼─────────────────────────┤
-│ REINFORCE             │ High       │ No           │ Quick baseline          │
-│ REINFORCE + baseline  │ Medium     │ Optional     │ Most use cases          │
-│ GRPO                  │ Low        │ No           │ Multi-sample rollouts   │
-│ RLOO                  │ Lowest     │ No           │ Small groups (K=2-4)    │
-└───────────────────────┴────────────┴──────────────┴─────────────────────────┘
-
-Recommended starting point: REINFORCE + baseline (simplest, good variance,
-no value head needed). If variance is still too high, switch to GRPO with K=4.
-
-For the SRE domain specifically:
-- Episodes are 100 steps long → significant credit assignment challenge
-- Rewards are dense (computed every step) → advantage normalization is key
-- Delayed effects (boot delay = 5 ticks) → GAE helps bridge the gap
-"""
-
-from __future__ import annotations
-
-import math
-from typing import List, Optional
-
-
-# ════════════════════════════════════════════════════════════════════════════════
-# Return / Advantage computation (shared across all loss functions)
-# ════════════════════════════════════════════════════════════════════════════════
-
-def compute_returns(
-    rewards: List[float],
-    gamma: float = 0.99,
-) -> List[float]:
-    """
-    Compute discounted returns (Monte Carlo) for each timestep.
-
-        G_t = r_t + γ * r_{t+1} + γ² * r_{t+2} + ... + γ^{T-t} * r_T
-
-    Args:
-        rewards: Per-step rewards [r_0, r_1, ..., r_{T-1}].
-        gamma: Discount factor. 0.99 = far-sighted, 0.9 = myopic.
-
-    Returns:
-        List of returns [G_0, G_1, ..., G_{T-1}] same length as rewards.
-    """
-    returns: List[float] = []
-    g = 0.0
-    for r in reversed(rewards):
-        g = r + gamma * g
-        returns.insert(0, g)
-    return returns
-
-
-def compute_gae(
-    rewards: List[float],
-    values: List[float],
-    gamma: float = 0.99,
-    lam: float = 0.95,
-) -> List[float]:
-    """
-    Generalized Advantage Estimation (GAE).
-
-        Â_t = Σ_{l=0}^{T-t-1} (γλ)^l δ_{t+l}
-
-    where δ_t = r_t + γ * V(s_{t+1}) - V(s_t) is the TD error.
-
-    GAE provides a bias-variance trade-off controlled by λ:
-        λ = 0  →  one-step TD (low variance, high bias)
-        λ = 1  →  Monte Carlo returns (high variance, no bias)
-
-    For SRE: λ=0.95 is a good default. The 5-tick boot delay means
-    actions have delayed effects — GAE with λ close to 1 helps propagate
-    credit across those gaps.
-
-    Args:
-        rewards: Per-step rewards [r_0, ..., r_{T-1}].
-        values:  State value estimates [V(s_0), ..., V(s_{T-1})].
-                 Pass a list of zeros for V=0 baseline (reduces to MC).
-        gamma:   Discount factor.
-        lam:     GAE lambda (trade-off parameter).
-
-    Returns:
-        List of GAE advantages [Â_0, ..., Â_{T-1}].
-    """
-    assert len(rewards) == len(values), f"len mismatch: rewards={len(rewards)} values={len(values)}"
-    advantages: List[float] = []
-    gae = 0.0
-    for t in reversed(range(len(rewards))):
-        if t == len(rewards) - 1:
-            next_value = 0.0  # Terminal state has V=0
-        else:
-            next_value = values[t + 1]
-        delta = rewards[t] + gamma * next_value - values[t]
-        gae = delta + gamma * lam * gae
-        advantages.insert(0, gae)
-    return advantages
-
-
-# ════════════════════════════════════════════════════════════════════════════════
-# Loss Functions
-# ════════════════════════════════════════════════════════════════════════════════
-
-def reinforce_loss(
-    log_probs: List[float],
-    returns: List[float],
-) -> float:
-    """
-    Vanilla REINFORCE (Williams, 1992).
-
-        L = -(1/T) Σ_t log π(a_t | s_t) · G_t
-
-    The gradient of this loss is an unbiased estimator of the policy gradient:
-        ∇J(θ) = E[Σ_t ∇log π(a_t|s_t) · G_t]
-
-    Properties:
-        - Unbiased but HIGH variance (no baseline)
-        - Simplest possible policy gradient method
-        - Good for initial prototyping, upgrade to baseline version ASAP
-
-    Args:
-        log_probs: log π(a_t | s_t) for each timestep.
-        returns:   Discounted returns G_t for each timestep.
-
-    Returns:
-        Scalar loss (minimize to maximize expected return).
-    """
-    assert len(log_probs) == len(returns), f"len mismatch: log_probs={len(log_probs)} returns={len(returns)}"
-    t = len(log_probs)
-    if t == 0:
-        return 0.0
-    return -sum(lp * g for lp, g in zip(log_probs, returns)) / t
-
-
-def reinforce_baseline_loss(
-    log_probs: List[float],
-    returns: List[float],
-    baselines: Optional[List[float]] = None,
-    normalize_advantage: bool = True,
-    eps: float = 1e-8,
-) -> float:
-    """
-    REINFORCE with baseline (variance reduction).
-
-        L = -(1/T) Σ_t log π(a_t | s_t) · Â_t
-
-    where Â_t = G_t - b_t is the advantage estimate.
-
-    The baseline b_t does NOT introduce bias (only reduces variance)
-    because E[∇log π(a|s) · b(s)] = 0 for any function b(s).
-
-    Baseline options:
-        - None (or zeros): Falls back to vanilla REINFORCE on returns
-        - Running mean of returns: Simple, effective, no extra model needed
-        - Learned value function: Most powerful, but needs value head
-
-    For SRE agents on Colab: use running-mean baseline (pass baselines as
-    the moving average of returns seen so far). No extra model needed.
-
-    Advantage normalization: Standardizes  to zero mean and unit variance.
-    This is critical for SRE because raw returns can span orders of magnitude
-    (0.001 vs 10.0) across episodes. Normalization keeps the learning rate
-    well-conditioned.
-
-    Args:
-        log_probs:         log π(a_t | s_t) for each timestep.
-        returns:           Discounted returns G_t for each timestep.
-        baselines:         Baseline estimates b_t. If None, uses mean(returns).
-        normalize_advantage: Whether to standardize advantages (recommended).
-        eps:               Epsilon for numerical stability in normalization.
-
-    Returns:
-        Scalar loss.
-    """
-    assert len(log_probs) == len(returns), f"len mismatch: log_probs={len(log_probs)} returns={len(returns)}"
-    t = len(log_probs)
-    if t == 0:
-        return 0.0
-
-    # Compute advantages
-    if baselines is None:
-        # Default baseline = mean of returns (simple but effective)
-        baseline_val = sum(returns) / len(returns)
-        advantages = [g - baseline_val for g in returns]
-    else:
-        assert len(baselines) == len(returns), f"len mismatch: baselines={len(baselines)} returns={len(returns)}"
-        advantages = [g - b for g, b in zip(returns, baselines)]
-
-    # Normalize advantages (critical for SRE reward scale)
-    if normalize_advantage and len(advantages) > 1:
-        mean_adv = sum(advantages) / len(advantages)
-        var_adv = sum((a - mean_adv) ** 2 for a in advantages) / len(advantages)
-        std_adv = math.sqrt(var_adv) + eps
-        advantages = [(a - mean_adv) / std_adv for a in advantages]
-
-    return -sum(lp * adv for lp, adv in zip(log_probs, advantages)) / t
-
-
-def grpo_loss(
-    log_probs_groups: List[List[float]],
-    rewards_groups: List[List[float]],
-    eps: float = 1e-8,
-) -> float:
-    """
-    Group Relative Policy Optimization (GRPO).
-
-    For each state s, generate K sample actions and compute group-relative
-    advantages without needing a value function:
-
-        Â_k = (r_k - μ_group) / σ_group
-
-        L = -(1/N) Σ_i (1/K_i) Σ_k log π(a_{i,k} | s_i) · Â_{i,k}
-
-    where μ_group and σ_group are the mean and std of rewards within group i.
-
-    This is the method used in DeepSeek-R1. It eliminates the need for a
-    value head entirely — the group statistics serve as the baseline.
-
-    Pros:
-        - No value function needed
-        - Low variance (group statistics absorb reward scale)
-        - Natural normalization
-
-    Cons:
-        - Requires K >= 2 rollouts per state (K * more compute)
-        - For K=1, falls back to REINFORCE (no baseline)
-
-    For SRE on Colab: Use K=4 with QLoRA on a T4. Each "group" is 4
-    different actions sampled for the same cluster state.
-
-    Args:
-        log_probs_groups: List of groups, each group is log π(a_k|s) for K samples.
-        rewards_groups:   List of groups, each group is reward_k for K samples.
-        eps:              Epsilon for std normalization.
-
-    Returns:
-        Scalar loss.
-    """
-    assert len(log_probs_groups) == len(rewards_groups), "group count mismatch"
-    if not log_probs_groups:
-        return 0.0
-
-    total_loss = 0.0
-    n_groups = 0
-
-    for log_probs, rewards in zip(log_probs_groups, rewards_groups):
-        assert len(log_probs) == len(rewards), f"group size mismatch: {len(log_probs)} vs {len(rewards)}"
-        k = len(log_probs)
-        if k == 0:
-            continue
-
-        # Group statistics
-        mean_r = sum(rewards) / k
-        var_r = sum((r - mean_r) ** 2 for r in rewards) / k
-        std_r = math.sqrt(var_r) + eps
-
-        # Normalized advantages
-        advantages = [(r - mean_r) / std_r for r in rewards]
-
-        # Policy gradient for this group
-        group_loss = -sum(lp * adv for lp, adv in zip(log_probs, advantages)) / k
-        total_loss += group_loss
-        n_groups += 1
-
-    return total_loss / max(1, n_groups)
-
-
-def rloo_loss(
-    log_probs_groups: List[List[float]],
-    rewards_groups: List[List[float]],
-) -> float:
-    """
-    REINFORCE Leave-One-Out (RLOO).
-
-    Similar to GRPO but uses a leave-one-out baseline instead of group
-    statistics. For each sample k in a group of K:
-
-        b_k = (1/(K-1)) Σ_{j≠k} r_j    (leave-one-out mean)
-
-        Â_k = r_k - b_k
-
-        L = -(1/N) Σ_i (1/K_i) Σ_k log π(a_{i,k} | s_i) · Â_{i,k}
-
-    RLOO has lower variance than GRPO for small group sizes (K=2-4) because
-    the baseline is computed from the actual other samples rather than a
-    statistical estimate. For K >= 8, GRPO and RLOO converge.
-
-    For SRE on Colab: Best choice when you can only afford K=2-4 rollouts.
-    The leave-one-out baseline is surprisingly effective.
-
-    Args:
-        log_probs_groups: List of groups, each group is log π(a_k|s) for K samples.
-        rewards_groups:   List of groups, each group is reward_k for K samples.
-
-    Returns:
-        Scalar loss.
-    """
-    assert len(log_probs_groups) == len(rewards_groups), "group count mismatch"
-    if not log_probs_groups:
-        return 0.0
-
-    total_loss = 0.0
-    n_groups = 0
-
-    for log_probs, rewards in zip(log_probs_groups, rewards_groups):
-        assert len(log_probs) == len(rewards), f"group size mismatch: {len(log_probs)} vs {len(rewards)}"
-        k = len(log_probs)
-        if k == 0:
-            continue
-
-        if k == 1:
-            # Single sample: no baseline possible, fall back to REINFORCE
-            total_loss += -log_probs[0] * rewards[0]
-            n_groups += 1
-            continue
-
-        # Leave-one-out baselines
-        sum_r = sum(rewards)
-        baselines = [(sum_r - r) / (k - 1) for r in rewards]
-        advantages = [r - b for r, b in zip(rewards, baselines)]
-
-        # Policy gradient
-        group_loss = -sum(lp * adv for lp, adv in zip(log_probs, advantages)) / k
-        total_loss += group_loss
-        n_groups += 1
-
-    return total_loss / max(1, n_groups)
-
-
-# ════════════════════════════════════════════════════════════════════════════════
-# Reward normalization utilities
-# ════════════════════════════════════════════════════════════════════════════════
-
-def normalize_rewards(
-    rewards: List[float],
-    running_mean: float = 0.0,
-    running_var: float = 1.0,
-    eps: float = 1e-8,
-) -> List[float]:
-    """
-    Normalize rewards using running statistics.
-
-    For SRE: Raw rewards are always negative (they're penalties). This function
-    shifts them to be centered around zero with unit variance, which is
-    essential for stable policy gradient updates.
-
-    On Colab: Maintain a running mean/var across episodes and pass them here.
-    Initialize with mean=0, var=1 and update with exponential moving average.
-
-        running_mean = 0.99 * running_mean + 0.01 * batch_mean
-        running_var  = 0.99 * running_var  + 0.01 * batch_var
-
-    Args:
-        rewards:       Raw rewards to normalize.
-        running_mean:  Running mean estimate across episodes.
-        running_var:   Running variance estimate across episodes.
-        eps:           Numerical stability constant.
-
-    Returns:
-        Normalized rewards (zero mean, unit variance relative to running stats).
-    """
-    std = math.sqrt(running_var) + eps
-    return [(r - running_mean) / std for r in rewards]
-
-
-def compute_reward_stats(rewards: List[float]) -> tuple[float, float]:
-    """
-    Compute mean and variance of a reward list.
-
-    Returns:
-        (mean, variance) tuple.
-    """
-    if not rewards:
-        return 0.0, 1.0
-    mean = sum(rewards) / len(rewards)
-    var = sum((r - mean) ** 2 for r in rewards) / len(rewards)
-    return mean, var
+"""
+AntiAtropos RL Loss Functions.
+
+Pure-Python implementations of policy gradient loss functions for LLM-based
+SRE agents. These are mathematically identical to their PyTorch counterparts
+and can be validated locally without GPU.
+
+When porting to PyTorch (Colab), simply replace:
+    - float ops with torch.tensor ops
+    - sum() with torch.sum()
+    - log() with torch.log()
+    - The gradient flows through log_prob automatically
+
+Loss function selection guide for the SRE domain:
+─────────────────────────────────────────────────
+┌───────────────────────┬────────────┬──────────────┬─────────────────────────┐
+│ Method                │ Variance   │ Needs Value  │ Best for                │
+├───────────────────────┼────────────┼──────────────┼─────────────────────────┤
+│ REINFORCE             │ High       │ No           │ Quick baseline          │
+│ REINFORCE + baseline  │ Medium     │ Optional     │ Most use cases          │
+│ GRPO                  │ Low        │ No           │ Multi-sample rollouts   │
+│ RLOO                  │ Lowest     │ No           │ Small groups (K=2-4)    │
+└───────────────────────┴────────────┴──────────────┴─────────────────────────┘
+
+Recommended starting point: REINFORCE + baseline (simplest, good variance,
+no value head needed). If variance is still too high, switch to GRPO with K=4.
+
+For the SRE domain specifically:
+- Episodes are 100 steps long → significant credit assignment challenge
+- Rewards are dense (computed every step) → advantage normalization is key
+- Delayed effects (boot delay = 5 ticks) → GAE helps bridge the gap
+"""
+
+from __future__ import annotations
+
+import math
+from typing import List, Optional
+
+
+# ════════════════════════════════════════════════════════════════════════════════
+# Return / Advantage computation (shared across all loss functions)
+# ════════════════════════════════════════════════════════════════════════════════
+
+def compute_returns(
+    rewards: List[float],
+    gamma: float = 0.99,
+) -> List[float]:
+    """
+    Compute discounted returns (Monte Carlo) for each timestep.
+
+        G_t = r_t + γ * r_{t+1} + γ² * r_{t+2} + ... + γ^{T-t} * r_T
+
+    Args:
+        rewards: Per-step rewards [r_0, r_1, ..., r_{T-1}].
+        gamma: Discount factor. 0.99 = far-sighted, 0.9 = myopic.
+
+    Returns:
+        List of returns [G_0, G_1, ..., G_{T-1}] same length as rewards.
+    """
+    returns: List[float] = []
+    g = 0.0
+    for r in reversed(rewards):
+        g = r + gamma * g
+        returns.insert(0, g)
+    return returns
+
+
+def compute_gae(
+    rewards: List[float],
+    values: List[float],
+    gamma: float = 0.99,
+    lam: float = 0.95,
+) -> List[float]:
+    """
+    Generalized Advantage Estimation (GAE).
+
+        Â_t = Σ_{l=0}^{T-t-1} (γλ)^l δ_{t+l}
+
+    where δ_t = r_t + γ * V(s_{t+1}) - V(s_t) is the TD error.
+
+    GAE provides a bias-variance trade-off controlled by λ:
+        λ = 0  →  one-step TD (low variance, high bias)
+        λ = 1  →  Monte Carlo returns (high variance, no bias)
+
+    For SRE: λ=0.95 is a good default. The 5-tick boot delay means
+    actions have delayed effects — GAE with λ close to 1 helps propagate
+    credit across those gaps.
+
+    Args:
+        rewards: Per-step rewards [r_0, ..., r_{T-1}].
+        values:  State value estimates [V(s_0), ..., V(s_{T-1})].
+                 Pass a list of zeros for V=0 baseline (reduces to MC).
+        gamma:   Discount factor.
+        lam:     GAE lambda (trade-off parameter).
+
+    Returns:
+        List of GAE advantages [Â_0, ..., Â_{T-1}].
+    """
+    assert len(rewards) == len(values), f"len mismatch: rewards={len(rewards)} values={len(values)}"
+    advantages: List[float] = []
+    gae = 0.0
+    for t in reversed(range(len(rewards))):
+        if t == len(rewards) - 1:
+            next_value = 0.0  # Terminal state has V=0
+        else:
+            next_value = values[t + 1]
+        delta = rewards[t] + gamma * next_value - values[t]
+        gae = delta + gamma * lam * gae
+        advantages.insert(0, gae)
+    return advantages
+
+
+# ════════════════════════════════════════════════════════════════════════════════
+# Loss Functions
+# ════════════════════════════════════════════════════════════════════════════════
+
+def reinforce_loss(
+    log_probs: List[float],
+    returns: List[float],
+) -> float:
+    """
+    Vanilla REINFORCE (Williams, 1992).
+
+        L = -(1/T) Σ_t log π(a_t | s_t) · G_t
+
+    The gradient of this loss is an unbiased estimator of the policy gradient:
+        ∇J(θ) = E[Σ_t ∇log π(a_t|s_t) · G_t]
+
+    Properties:
+        - Unbiased but HIGH variance (no baseline)
+        - Simplest possible policy gradient method
+        - Good for initial prototyping, upgrade to baseline version ASAP
+
+    Args:
+        log_probs: log π(a_t | s_t) for each timestep.
+        returns:   Discounted returns G_t for each timestep.
+
+    Returns:
+        Scalar loss (minimize to maximize expected return).
+    """
+    assert len(log_probs) == len(returns), f"len mismatch: log_probs={len(log_probs)} returns={len(returns)}"
+    t = len(log_probs)
+    if t == 0:
+        return 0.0
+    return -sum(lp * g for lp, g in zip(log_probs, returns)) / t
+
+
+def reinforce_baseline_loss(
+    log_probs: List[float],
+    returns: List[float],
+    baselines: Optional[List[float]] = None,
+    normalize_advantage: bool = True,
+    eps: float = 1e-8,
+) -> float:
+    """
+    REINFORCE with baseline (variance reduction).
+
+        L = -(1/T) Σ_t log π(a_t | s_t) · Â_t
+
+    where Â_t = G_t - b_t is the advantage estimate.
+
+    The baseline b_t does NOT introduce bias (only reduces variance)
+    because E[∇log π(a|s) · b(s)] = 0 for any function b(s).
+
+    Baseline options:
+        - None (or zeros): Falls back to vanilla REINFORCE on returns
+        - Running mean of returns: Simple, effective, no extra model needed
+        - Learned value function: Most powerful, but needs value head
+
+    For SRE agents on Colab: use running-mean baseline (pass baselines as
+    the moving average of returns seen so far). No extra model needed.
+
+    Advantage normalization: Standardizes  to zero mean and unit variance.
+    This is critical for SRE because raw returns can span orders of magnitude
+    (0.001 vs 10.0) across episodes. Normalization keeps the learning rate
+    well-conditioned.
+
+    Args:
+        log_probs:         log π(a_t | s_t) for each timestep.
+        returns:           Discounted returns G_t for each timestep.
+        baselines:         Baseline estimates b_t. If None, uses mean(returns).
+        normalize_advantage: Whether to standardize advantages (recommended).
+        eps:               Epsilon for numerical stability in normalization.
+
+    Returns:
+        Scalar loss.
+    """
+    assert len(log_probs) == len(returns), f"len mismatch: log_probs={len(log_probs)} returns={len(returns)}"
+    t = len(log_probs)
+    if t == 0:
+        return 0.0
+
+    # Compute advantages
+    if baselines is None:
+        # Default baseline = mean of returns (simple but effective)
+        baseline_val = sum(returns) / len(returns)
+        advantages = [g - baseline_val for g in returns]
+    else:
+        assert len(baselines) == len(returns), f"len mismatch: baselines={len(baselines)} returns={len(returns)}"
+        advantages = [g - b for g, b in zip(returns, baselines)]
+
+    # Normalize advantages (critical for SRE reward scale)
+    if normalize_advantage and len(advantages) > 1:
+        mean_adv = sum(advantages) / len(advantages)
+        var_adv = sum((a - mean_adv) ** 2 for a in advantages) / len(advantages)
+        std_adv = math.sqrt(var_adv) + eps
+        advantages = [(a - mean_adv) / std_adv for a in advantages]
+
+    return -sum(lp * adv for lp, adv in zip(log_probs, advantages)) / t
+
+
+def grpo_loss(
+    log_probs_groups: List[List[float]],
+    rewards_groups: List[List[float]],
+    eps: float = 1e-8,
+) -> float:
+    """
+    Group Relative Policy Optimization (GRPO).
+
+    For each state s, generate K sample actions and compute group-relative
+    advantages without needing a value function:
+
+        Â_k = (r_k - μ_group) / σ_group
+
+        L = -(1/N) Σ_i (1/K_i) Σ_k log π(a_{i,k} | s_i) · Â_{i,k}
+
+    where μ_group and σ_group are the mean and std of rewards within group i.
+
+    This is the method used in DeepSeek-R1. It eliminates the need for a
+    value head entirely — the group statistics serve as the baseline.
+
+    Pros:
+        - No value function needed
+        - Low variance (group statistics absorb reward scale)
+        - Natural normalization
+
+    Cons:
+        - Requires K >= 2 rollouts per state (K * more compute)
+        - For K=1, falls back to REINFORCE (no baseline)
+
+    For SRE on Colab: Use K=4 with QLoRA on a T4. Each "group" is 4
+    different actions sampled for the same cluster state.
+
+    Args:
+        log_probs_groups: List of groups, each group is log π(a_k|s) for K samples.
+        rewards_groups:   List of groups, each group is reward_k for K samples.
+        eps:              Epsilon for std normalization.
+
+    Returns:
+        Scalar loss.
+    """
+    assert len(log_probs_groups) == len(rewards_groups), "group count mismatch"
+    if not log_probs_groups:
+        return 0.0
+
+    total_loss = 0.0
+    n_groups = 0
+
+    for log_probs, rewards in zip(log_probs_groups, rewards_groups):
+        assert len(log_probs) == len(rewards), f"group size mismatch: {len(log_probs)} vs {len(rewards)}"
+        k = len(log_probs)
+        if k == 0:
+            continue
+
+        # Group statistics
+        mean_r = sum(rewards) / k
+        var_r = sum((r - mean_r) ** 2 for r in rewards) / k
+        std_r = math.sqrt(var_r) + eps
+
+        # Normalized advantages
+        advantages = [(r - mean_r) / std_r for r in rewards]
+
+        # Policy gradient for this group
+        group_loss = -sum(lp * adv for lp, adv in zip(log_probs, advantages)) / k
+        total_loss += group_loss
+        n_groups += 1
+
+    return total_loss / max(1, n_groups)
+
+
+def rloo_loss(
+    log_probs_groups: List[List[float]],
+    rewards_groups: List[List[float]],
+) -> float:
+    """
+    REINFORCE Leave-One-Out (RLOO).
+
+    Similar to GRPO but uses a leave-one-out baseline instead of group
+    statistics. For each sample k in a group of K:
+
+        b_k = (1/(K-1)) Σ_{j≠k} r_j    (leave-one-out mean)
+
+        Â_k = r_k - b_k
+
+        L = -(1/N) Σ_i (1/K_i) Σ_k log π(a_{i,k} | s_i) · Â_{i,k}
+
+    RLOO has lower variance than GRPO for small group sizes (K=2-4) because
+    the baseline is computed from the actual other samples rather than a
+    statistical estimate. For K >= 8, GRPO and RLOO converge.
+
+    For SRE on Colab: Best choice when you can only afford K=2-4 rollouts.
+    The leave-one-out baseline is surprisingly effective.
+
+    Args:
+        log_probs_groups: List of groups, each group is log π(a_k|s) for K samples.
+        rewards_groups:   List of groups, each group is reward_k for K samples.
+
+    Returns:
+        Scalar loss.
+    """
+    assert len(log_probs_groups) == len(rewards_groups), "group count mismatch"
+    if not log_probs_groups:
+        return 0.0
+
+    total_loss = 0.0
+    n_groups = 0
+
+    for log_probs, rewards in zip(log_probs_groups, rewards_groups):
+        assert len(log_probs) == len(rewards), f"group size mismatch: {len(log_probs)} vs {len(rewards)}"
+        k = len(log_probs)
+        if k == 0:
+            continue
+
+        if k == 1:
+            # Single sample: no baseline possible, fall back to REINFORCE
+            total_loss += -log_probs[0] * rewards[0]
+            n_groups += 1
+            continue
+
+        # Leave-one-out baselines
+        sum_r = sum(rewards)
+        baselines = [(sum_r - r) / (k - 1) for r in rewards]
+        advantages = [r - b for r, b in zip(rewards, baselines)]
+
+        # Policy gradient
+        group_loss = -sum(lp * adv for lp, adv in zip(log_probs, advantages)) / k
+        total_loss += group_loss
+        n_groups += 1
+
+    return total_loss / max(1, n_groups)
+
+
+# ════════════════════════════════════════════════════════════════════════════════
+# Reward normalization utilities
+# ════════════════════════════════════════════════════════════════════════════════
+
+def normalize_rewards(
+    rewards: List[float],
+    running_mean: float = 0.0,
+    running_var: float = 1.0,
+    eps: float = 1e-8,
+) -> List[float]:
+    """
+    Normalize rewards using running statistics.
+
+    For SRE: Raw rewards are always negative (they're penalties). This function
+    shifts them to be centered around zero with unit variance, which is
+    essential for stable policy gradient updates.
+
+    On Colab: Maintain a running mean/var across episodes and pass them here.
+    Initialize with mean=0, var=1 and update with exponential moving average.
+
+        running_mean = 0.99 * running_mean + 0.01 * batch_mean
+        running_var  = 0.99 * running_var  + 0.01 * batch_var
+
+    Args:
+        rewards:       Raw rewards to normalize.
+        running_mean:  Running mean estimate across episodes.
+        running_var:   Running variance estimate across episodes.
+        eps:           Numerical stability constant.
+
+    Returns:
+        Normalized rewards (zero mean, unit variance relative to running stats).
+    """
+    std = math.sqrt(running_var) + eps
+    return [(r - running_mean) / std for r in rewards]
+
+
+def compute_reward_stats(rewards: List[float]) -> tuple[float, float]:
+    """
+    Compute mean and variance of a reward list.
+
+    Returns:
+        (mean, variance) tuple.
+    """
+    if not rewards:
+        return 0.0, 1.0
+    mean = sum(rewards) / len(rewards)
+    var = sum((r - mean) ** 2 for r in rewards) / len(rewards)
+    return mean, var

training/trainer.py

@@ -1,559 +1,559 @@
-"""
-AntiAtropos Training Loop.
-
-Orchestrates episode collection, reward computation, and loss calculation
-for training LLM-based SRE agents. Works with the local simulator in
-pure-Python mode (no AWS/GPU needed for validation).
-
-On Colab: Replace EpisodeCollector's "model" with a real QLoRA-backed
-transformers model. The rest of the pipeline stays the same.
-"""
-
-from __future__ import annotations
-
-import random
-import math
-from dataclasses import dataclass, field
-from typing import List, Optional, Protocol, Callable
-
-import sys
-import os
-sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
-
-from simulator import ClusterSimulator, NodeStatus, COST_PER_CAPACITY_UNIT_PER_HOUR
-from stability import (
-    compute_lyapunov, compute_reward, compute_barrier,
-    normalize_reward, smooth_sla_penalty,
-)
-from .losses import (
-    compute_returns, compute_gae,
-    reinforce_loss, reinforce_baseline_loss,
-    grpo_loss, rloo_loss,
-    normalize_rewards, compute_reward_stats,
-)
-
-
-# ════════════════════════════════════════════════════════════════════════════════
-# Configuration
-# ════════════════════════════════════════════════════════════════════════════════
-
-LOSS_REINFORCE = "reinforce"
-LOSS_REINFORCE_BASELINE = "reinforce_baseline"
-LOSS_GRPO = "grpo"
-LOSS_RLOO = "rloo"
-
-VALID_LOSSES = {LOSS_REINFORCE, LOSS_REINFORCE_BASELINE, LOSS_GRPO, LOSS_RLOO}
-
-
-@dataclass
-class TrainingConfig:
-    """Configuration for the SRE training loop."""
-
-    # Episode settings
-    n_nodes: int = 5
-    max_steps: int = 100
-    tasks: List[str] = field(default_factory=lambda: ["task-1", "task-2", "task-3"])
-
-    # Loss function
-    loss_fn: str = LOSS_REINFORCE_BASELINE  # Recommended starting point
-    gamma: float = 0.99          # Discount factor
-    gae_lambda: float = 0.95     # GAE lambda (only used with GAE advantages)
-
-    # GRPO / RLOO settings
-    n_samples_per_state: int = 4  # K rollouts per state for GRPO/RLOO
-
-    # Reward normalization
-    normalize_rewards: bool = True
-    reward_ema_alpha: float = 0.01  # Exponential moving average update rate
-
-    # Advantage normalization
-    normalize_advantages: bool = True  # Standardize advantages (critical for SRE)
-
-    # Logging
-    log_every: int = 10  # Log every N episodes
-
-
-# ════════════════════════════════════════════════════════════════════════════════
-# Model Protocol (abstraction for real LLM or mock)
-# ════════════════════════════════════════════════════════════════════════════════
-
-class PolicyModel(Protocol):
-    """Interface that both real LLMs and mock models must implement."""
-
-    def get_log_prob(self, prompt: str, action_text: str) -> float:
-        """Return log π(action_text | prompt) under the current policy."""
-        ...
-
-    def generate(self, prompt: str) -> str:
-        """Sample an action from the current policy."""
-        ...
-
-
-class MockPolicyModel:
-    """
-    Random policy for local testing. Generates random valid actions
-    and returns stochastic log probabilities.
-
-    This is NOT for training — it's for validating the training pipeline
-    (episode collection, reward computation, loss calculation) end-to-end
-    before connecting a real model.
-
-    Unlike a truly uniform policy (which would produce zero gradient with
-    advantage normalization), this model returns varying log probs per
-    action, simulating a real LLM that prefers some actions over others.
-    This ensures the training pipeline produces non-trivial losses.
-    """
-
-    def __init__(self, n_nodes: int = 5, seed: int = 42):
-        self._rng = random.Random(seed)
-        self._n_nodes = n_nodes
-        # Action-dependent log probs: different action types get different
-        # log probs (simulating a real model that has preferences).
-        # Base log prob ~ log(1/50) ≈ -3.9, with per-action noise.
-        self._n_choices = 5 * n_nodes
-        self._base_log_prob = math.log(1.0 / self._n_choices)
-
-    def get_log_prob(self, prompt: str, action_text: str) -> float:
-        """Return stochastic log probability (varies per action)."""
-        # Add Gaussian noise to simulate a real model's varying confidence.
-        # std=0.5 produces meaningful variation while staying in a plausible
-        # range for LLM token log-probs.
-        noise = self._rng.gauss(0, 0.5)
-        return self._base_log_prob + noise
-
-    def generate(self, prompt: str) -> str:
-        """Generate a random valid action as JSON string."""
-        import json
-        action_types = ["SCALE_UP", "SCALE_DOWN", "REROUTE_TRAFFIC", "SHED_LOAD", "NO_OP"]
-        node_id = f"node-{self._rng.randint(0, self._n_nodes - 1)}"
-        action_type = self._rng.choice(action_types)
-        parameter = round(self._rng.random(), 2)
-        return json.dumps({
-            "action_type": action_type,
-            "target_node_id": node_id,
-            "parameter": parameter,
-        })
-
-
-# ════════════════════════════════════════════════════════════════════════════════
-# Observation formatting (mirrors inference.py logic)
-# ════════════════════════════════════════════════════════════════════════════════
-
-MAX_QUEUE_NORM = 200.0
-MAX_LATENCY_NORM = 1000.0
-MAX_REQUEST_RATE_NORM = 100.0
-ALPHA, BETA, GAMMA, DELTA = 0.002, 0.01, 10.0, 0.005
-
-
-def format_observation(nodes: List[dict], task_id: str, step: int, max_steps: int) -> str:
-    """
-    Format simulator state as a text prompt for the model.
-
-    This mirrors inference.py's build_user_prompt and observation_for_model.
-    """
-    import json
-    node_data = []
-    for n in nodes:
-        node_data.append({
-            "node_id": n["node_id"],
-            "status": n["status"] if isinstance(n["status"], str) else n["status"].value,
-            "is_vip": n.get("is_vip", False),
-            "queue_depth": min(1.0, max(0.0, n["queue_depth"] / MAX_QUEUE_NORM)),
-            "latency_ms": min(1.0, max(0.0, n["latency_ms"] / MAX_LATENCY_NORM)),
-            "cpu_utilization": min(1.0, max(0.0, n.get("cpu_utilization", 0.0))),
-            "incoming_request_rate": min(1.0, max(0.0, n["incoming_request_rate"] / MAX_REQUEST_RATE_NORM)),
-        })
-    obs = {"task_id": task_id, "step": step, "max_steps": max_steps, "nodes": node_data}
-    return json.dumps(obs, separators=(",", ":"))
-
-
-def parse_action(action_text: str) -> dict:
-    """Parse model output into an action dict."""
-    import json
-    try:
-        data = json.loads(action_text)
-        return {
-            "action_type": str(data.get("action_type", "NO_OP")).upper(),
-            "target_node_id": str(data.get("target_node_id", "node-0")),
-            "parameter": float(data.get("parameter", 0.0)),
-        }
-    except (json.JSONDecodeError, ValueError):
-        return {"action_type": "NO_OP", "target_node_id": "node-0", "parameter": 0.0}
-
-
-# ════════════════════════════════════════════════════════════════════════════════
-# Episode Collection
-# ════════════════════════════════════════════════════════════════════════════════
-
-@dataclass
-class StepRecord:
-    """A single step in an episode."""
-    prompt: str           # Observation formatted as text
-    action_text: str      # Model output (JSON string)
-    log_prob: float       # log π(action | prompt)
-    reward: float         # Raw reward for this step
-    reward_normalized: float  # Normalized reward
-
-
-@dataclass
-class EpisodeRecord:
-    """A complete episode trajectory."""
-    task_id: str
-    steps: List[StepRecord] = field(default_factory=list)
-    total_reward: float = 0.0
-    total_reward_normalized: float = 0.0
-    avg_reward: float = 0.0
-    sla_violations: int = 0
-    final_lyapunov: float = 0.0
-
-
-class EpisodeCollector:
-    """
-    Collects episodes by running the simulator with a policy model.
-
-    This is the bridge between the simulator (physics) and the training
-    pipeline (loss computation). It produces EpisodeRecords that feed
-    directly into the loss functions.
-    """
-
-    def __init__(self, config: TrainingConfig):
-        self._config = config
-        self._sim = ClusterSimulator(n_nodes=config.n_nodes)
-
-    def collect_episode(
-        self,
-        model: PolicyModel,
-        task_id: str,
-        seed: Optional[int] = None,
-    ) -> EpisodeRecord:
-        """Run one episode and collect step-level data."""
-        cfg = self._config
-        self._sim.reset(task_id=task_id, seed=seed)
-
-        record = EpisodeRecord(task_id=task_id)
-        prev_lyapunov = 0.0
-
-        for step in range(1, cfg.max_steps + 1):
-            # 1. Get observation
-            nodes_true = self._sim.state(for_agent=False)
-            nodes_obs = self._sim.state(for_agent=True)
-            prompt = format_observation(nodes_obs, task_id, step, cfg.max_steps)
-
-            # 2. Get action from model
-            action_text = model.generate(prompt)
-            log_prob = model.get_log_prob(prompt, action_text)
-
-            # 3. Apply action
-            action = parse_action(action_text)
-            class _A:
-                pass
-            a = _A()
-            a.action_type = action["action_type"]
-            a.target_node_id = action["target_node_id"]
-            a.parameter = action["parameter"]
-            self._sim.apply_action(a)
-
-            # 4. Tick
-            self._sim.tick()
-
-            # 5. Compute reward (mirrors environment.py)
-            nodes_true = self._sim.state(for_agent=False)
-            current_lyapunov = compute_lyapunov(nodes_true)
-
-            # Importance-weighted average latency
-            w_lat = 0.0
-            w_sum = 0.0
-            for n in nodes_true:
-                w = n.get("importance_weight", 1.0)
-                lat = MAX_LATENCY_NORM if n["status"] == NodeStatus.FAILED else n["latency_ms"]
-                w_lat += w * lat
-                w_sum += w
-            avg_lat_norm = min(1.0, max(0.0, (w_lat / w_sum / MAX_LATENCY_NORM) if w_sum > 0 else 1.0))
-
-            # Error rate
-            total_in = sum(n.get("incoming_request_rate", 0) * n.get("importance_weight", 1.0) for n in nodes_true)
-            total_drop = sum(n.get("dropped_requests", 0) * n.get("importance_weight", 1.0) for n in nodes_true)
-            error_rate = min(1.0, total_drop / total_in) if total_in > 0 else 0.0
-
-            sla_step = smooth_sla_penalty(avg_lat_norm, error_rate)
-            if avg_lat_norm > 0.20 or error_rate > 0.05:
-                record.sla_violations += 1
-
-            # Cost
-            total_cap = 0
-            for n in nodes_true:
-                if n["status"] != NodeStatus.FAILED:
-                    total_cap += int(n.get("capacity_units", 0)) + int(n.get("pending_capacity_units", 0))
-            cost = total_cap * COST_PER_CAPACITY_UNIT_PER_HOUR
-
-            barrier = compute_barrier(nodes_true)
-            raw_reward = compute_reward(
-                prev_lyapunov, current_lyapunov, cost, sla_step,
-                ALPHA, BETA, GAMMA, barrier, DELTA,
-            )
-            norm_reward = normalize_reward(raw_reward)
-
-            record.steps.append(StepRecord(
-                prompt=prompt,
-                action_text=action_text,
-                log_prob=log_prob,
-                reward=raw_reward,
-                reward_normalized=norm_reward,
-            ))
-            record.total_reward += raw_reward
-            record.total_reward_normalized += norm_reward
-            prev_lyapunov = current_lyapunov
-
-        record.avg_reward = record.total_reward / max(1, len(record.steps))
-        record.final_lyapunov = prev_lyapunov
-        return record
-
-    def collect_group(
-        self,
-        model: PolicyModel,
-        task_id: str,
-        k: int,
-        seed: Optional[int] = None,
-    ) -> List[EpisodeRecord]:
-        """
-        Collect K episodes from the same initial state (for GRPO/RLOO).
-
-        Uses the same seed for all K episodes so they start from the same
-        domain randomization, but different model samples produce different
-        trajectories.
-        """
-        return [self.collect_episode(model, task_id, seed=seed) for _ in range(k)]
-
-
-# ════════════════════════════════════════════════════════════════════════════════
-# Trainer
-# ════════════════════════════════════════════════════════════════════════════════
-
-class SRETrainer:
-    """
-    Main training orchestrator for AntiAtropos SRE agents.
-
-    Usage (local validation with MockPolicyModel):
-        config = TrainingConfig(loss_fn="reinforce_baseline")
-        trainer = SRETrainer(config)
-        model = MockPolicyModel()
-        metrics = trainer.train_step(model, task_id="task-1", seed=42)
-
-    Usage (Colab with real model):
-        config = TrainingConfig(loss_fn="grpo", n_samples_per_state=4)
-        trainer = SRETrainer(config)
-        model = QLoRAModel(...)  # Your transformers model
-        for epoch in range(num_epochs):
-            for task in config.tasks:
-                metrics = trainer.train_step(model, task_id=task)
-                model.update(metrics["loss"])  # Backprop
-    """
-
-    def __init__(self, config: TrainingConfig):
-        assert config.loss_fn in VALID_LOSSES, f"Unknown loss: {config.loss_fn}"
-        self._config = config
-        self._collector = EpisodeCollector(config)
-        self._running_reward_mean = 0.0
-        self._running_reward_var = 1.0
-        self._episode_count = 0
-
-    def train_step(
-        self,
-        model: PolicyModel,
-        task_id: str,
-        seed: Optional[int] = None,
-    ) -> dict:
-        """
-        Execute one training step: collect episode(s) → compute loss.
-
-        Returns a metrics dict with:
-            - loss: The computed loss value
-            - avg_reward: Average raw reward across the episode
-            - avg_norm_reward: Average normalized reward
-            - episode_length: Number of steps
-            - sla_violations: Number of SLA violations
-            - final_lyapunov: Lyapunov energy at episode end
-            - reward_mean/var: Running reward statistics
-        """
-        cfg = self._config
-
-        if cfg.loss_fn in (LOSS_GRPO, LOSS_RLOO):
-            return self._train_step_grouped(model, task_id, seed)
-        else:
-            return self._train_step_single(model, task_id, seed)
-
-    def _train_step_single(
-        self,
-        model: PolicyModel,
-        task_id: str,
-        seed: Optional[int] = None,
-    ) -> dict:
-        """Train step for REINFORCE / REINFORCE+baseline."""
-        cfg = self._config
-
-        # 1. Collect episode
-        episode = self._collector.collect_episode(model, task_id, seed=seed)
-
-        # 2. Extract rewards and log probs
-        rewards = [s.reward for s in episode.steps]
-        log_probs = [s.log_prob for s in episode.steps]
-
-        # 3. Update running reward stats
-        ep_mean, ep_var = compute_reward_stats(rewards)
-        self._running_reward_mean = (
-            (1 - cfg.reward_ema_alpha) * self._running_reward_mean
-            + cfg.reward_ema_alpha * ep_mean
-        )
-        self._running_reward_var = (
-            (1 - cfg.reward_ema_alpha) * self._running_reward_var
-            + cfg.reward_ema_alpha * ep_var
-        )
-
-        # 4. Optionally normalize rewards
-        if cfg.normalize_rewards:
-            rewards = normalize_rewards(
-                rewards, self._running_reward_mean, self._running_reward_var
-            )
-
-        # 5. Compute returns
-        returns = compute_returns(rewards, gamma=cfg.gamma)
-
-        # 6. Compute loss
-        if cfg.loss_fn == LOSS_REINFORCE:
-            loss = reinforce_loss(log_probs, returns)
-        elif cfg.loss_fn == LOSS_REINFORCE_BASELINE:
-            # Use running mean as baseline
-            baselines = [self._running_reward_mean] * len(returns)
-            loss = reinforce_baseline_loss(
-                log_probs, returns, baselines,
-                normalize_advantage=cfg.normalize_advantages,
-            )
-        else:
-            raise ValueError(f"Unexpected loss_fn: {cfg.loss_fn}")
-
-        self._episode_count += 1
-
-        return {
-            "loss": loss,
-            "avg_reward": episode.avg_reward,
-            "avg_norm_reward": episode.total_reward_normalized / max(1, len(episode.steps)),
-            "episode_length": len(episode.steps),
-            "sla_violations": episode.sla_violations,
-            "final_lyapunov": episode.final_lyapunov,
-            "reward_mean": self._running_reward_mean,
-            "reward_var": self._running_reward_var,
-            "task_id": task_id,
-            "episode": episode,
-        }
-
-    def _train_step_grouped(
-        self,
-        model: PolicyModel,
-        task_id: str,
-        seed: Optional[int] = None,
-    ) -> dict:
-        """Train step for GRPO / RLOO."""
-        cfg = self._config
-        k = cfg.n_samples_per_state
-
-        # 1. Collect K episodes (same seed → same domain randomization)
-        episodes = self._collector.collect_group(model, task_id, k=k, seed=seed)
-
-        # 2. For each step position, form groups across episodes
-        #    (assumes all episodes have same length)
-        min_len = min(len(ep.steps) for ep in episodes)
-
-        log_probs_groups = []
-        rewards_groups = []
-
-        for t in range(min_len):
-            step_lps = []
-            step_rs = []
-            for ep in episodes:
-                step_lps.append(ep.steps[t].log_prob)
-                step_rs.append(ep.steps[t].reward)
-            log_probs_groups.append(step_lps)
-            rewards_groups.append(step_rs)
-
-        # 3. Update running stats
-        all_rewards = [s.reward for ep in episodes for s in ep.steps]
-        ep_mean, ep_var = compute_reward_stats(all_rewards)
-        self._running_reward_mean = (
-            (1 - cfg.reward_ema_alpha) * self._running_reward_mean
-            + cfg.reward_ema_alpha * ep_mean
-        )
-        self._running_reward_var = (
-            (1 - cfg.reward_ema_alpha) * self._running_reward_var
-            + cfg.reward_ema_alpha * ep_var
-        )
-
-        # 4. Normalize rewards
-        if cfg.normalize_rewards:
-            rewards_groups = [
-                normalize_rewards(rs, self._running_reward_mean, self._running_reward_var)
-                for rs in rewards_groups
-            ]
-
-        # 5. Compute loss
-        if cfg.loss_fn == LOSS_GRPO:
-            loss = grpo_loss(log_probs_groups, rewards_groups)
-        elif cfg.loss_fn == LOSS_RLOO:
-            loss = rloo_loss(log_probs_groups, rewards_groups)
-        else:
-            raise ValueError(f"Unexpected grouped loss_fn: {cfg.loss_fn}")
-
-        # 6. Aggregate metrics across episodes
-        avg_reward = sum(ep.avg_reward for ep in episodes) / len(episodes)
-        avg_norm = sum(
-            ep.total_reward_normalized / max(1, len(ep.steps)) for ep in episodes
-        ) / len(episodes)
-        total_sla = sum(ep.sla_violations for ep in episodes)
-        avg_lyapunov = sum(ep.final_lyapunov for ep in episodes) / len(episodes)
-
-        self._episode_count += k
-
-        return {
-            "loss": loss,
-            "avg_reward": avg_reward,
-            "avg_norm_reward": avg_norm,
-            "episode_length": min_len,
-            "sla_violations": total_sla,
-            "final_lyapunov": avg_lyapunov,
-            "reward_mean": self._running_reward_mean,
-            "reward_var": self._running_reward_var,
-            "task_id": task_id,
-            "episodes": episodes,
-        }
-
-    def train_epoch(
-        self,
-        model: PolicyModel,
-        seed: Optional[int] = None,
-    ) -> List[dict]:
-        """
-        Run one training step per task in the curriculum.
-
-        Returns a list of metrics dicts (one per task).
-        """
-        results = []
-        for task_id in self._config.tasks:
-            step_seed = seed + hash(task_id) % 1000 if seed is not None else None
-            metrics = self.train_step(model, task_id, seed=step_seed)
-            results.append(metrics)
-            if self._episode_count % self._config.log_every == 0:
-                self._log_metrics(metrics)
-        return results
-
-    def _log_metrics(self, metrics: dict) -> None:
-        """Print training metrics."""
-        print(
-            f"[Episode {self._episode_count}] "
-            f"task={metrics['task_id']} "
-            f"loss={metrics['loss']:.4f} "
-            f"avg_reward={metrics['avg_reward']:.4f} "
-            f"avg_norm_reward={metrics['avg_norm_reward']:.4f} "
-            f"sla_violations={metrics['sla_violations']} "
-            f"lyapunov={metrics['final_lyapunov']:.1f} "
-            f"reward_mean={metrics['reward_mean']:.4f} "
-            f"reward_var={metrics['reward_var']:.4f}"
-        )
+"""
+AntiAtropos Training Loop.
+
+Orchestrates episode collection, reward computation, and loss calculation
+for training LLM-based SRE agents. Works with the local simulator in
+pure-Python mode (no AWS/GPU needed for validation).
+
+On Colab: Replace EpisodeCollector's "model" with a real QLoRA-backed
+transformers model. The rest of the pipeline stays the same.
+"""
+
+from __future__ import annotations
+
+import random
+import math
+from dataclasses import dataclass, field
+from typing import List, Optional, Protocol, Callable
+
+import sys
+import os
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
+
+from simulator import ClusterSimulator, NodeStatus, COST_PER_CAPACITY_UNIT_PER_HOUR
+from stability import (
+    compute_lyapunov, compute_reward, compute_barrier,
+    normalize_reward, smooth_sla_penalty,
+)
+from .losses import (
+    compute_returns, compute_gae,
+    reinforce_loss, reinforce_baseline_loss,
+    grpo_loss, rloo_loss,
+    normalize_rewards, compute_reward_stats,
+)
+
+
+# ════════════════════════════════════════════════════════════════════════════════
+# Configuration
+# ════════════════════════════════════════════════════════════════════════════════
+
+LOSS_REINFORCE = "reinforce"
+LOSS_REINFORCE_BASELINE = "reinforce_baseline"
+LOSS_GRPO = "grpo"
+LOSS_RLOO = "rloo"
+
+VALID_LOSSES = {LOSS_REINFORCE, LOSS_REINFORCE_BASELINE, LOSS_GRPO, LOSS_RLOO}
+
+
+@dataclass
+class TrainingConfig:
+    """Configuration for the SRE training loop."""
+
+    # Episode settings
+    n_nodes: int = 5
+    max_steps: int = 100
+    tasks: List[str] = field(default_factory=lambda: ["task-1", "task-2", "task-3"])
+
+    # Loss function
+    loss_fn: str = LOSS_REINFORCE_BASELINE  # Recommended starting point
+    gamma: float = 0.99          # Discount factor
+    gae_lambda: float = 0.95     # GAE lambda (only used with GAE advantages)
+
+    # GRPO / RLOO settings
+    n_samples_per_state: int = 4  # K rollouts per state for GRPO/RLOO
+
+    # Reward normalization
+    normalize_rewards: bool = True
+    reward_ema_alpha: float = 0.01  # Exponential moving average update rate
+
+    # Advantage normalization
+    normalize_advantages: bool = True  # Standardize advantages (critical for SRE)
+
+    # Logging
+    log_every: int = 10  # Log every N episodes
+
+
+# ════════════════════════════════════════════════════════════════════════════════
+# Model Protocol (abstraction for real LLM or mock)
+# ════════════════════════════════════════════════════════════════════════════════
+
+class PolicyModel(Protocol):
+    """Interface that both real LLMs and mock models must implement."""
+
+    def get_log_prob(self, prompt: str, action_text: str) -> float:
+        """Return log π(action_text | prompt) under the current policy."""
+        ...
+
+    def generate(self, prompt: str) -> str:
+        """Sample an action from the current policy."""
+        ...
+
+
+class MockPolicyModel:
+    """
+    Random policy for local testing. Generates random valid actions
+    and returns stochastic log probabilities.
+
+    This is NOT for training — it's for validating the training pipeline
+    (episode collection, reward computation, loss calculation) end-to-end
+    before connecting a real model.
+
+    Unlike a truly uniform policy (which would produce zero gradient with
+    advantage normalization), this model returns varying log probs per
+    action, simulating a real LLM that prefers some actions over others.
+    This ensures the training pipeline produces non-trivial losses.
+    """
+
+    def __init__(self, n_nodes: int = 5, seed: int = 42):
+        self._rng = random.Random(seed)
+        self._n_nodes = n_nodes
+        # Action-dependent log probs: different action types get different
+        # log probs (simulating a real model that has preferences).
+        # Base log prob ~ log(1/50) ≈ -3.9, with per-action noise.
+        self._n_choices = 5 * n_nodes
+        self._base_log_prob = math.log(1.0 / self._n_choices)
+
+    def get_log_prob(self, prompt: str, action_text: str) -> float:
+        """Return stochastic log probability (varies per action)."""
+        # Add Gaussian noise to simulate a real model's varying confidence.
+        # std=0.5 produces meaningful variation while staying in a plausible
+        # range for LLM token log-probs.
+        noise = self._rng.gauss(0, 0.5)
+        return self._base_log_prob + noise
+
+    def generate(self, prompt: str) -> str:
+        """Generate a random valid action as JSON string."""
+        import json
+        action_types = ["SCALE_UP", "SCALE_DOWN", "REROUTE_TRAFFIC", "SHED_LOAD", "NO_OP"]
+        node_id = f"node-{self._rng.randint(0, self._n_nodes - 1)}"
+        action_type = self._rng.choice(action_types)
+        parameter = round(self._rng.random(), 2)
+        return json.dumps({
+            "action_type": action_type,
+            "target_node_id": node_id,
+            "parameter": parameter,
+        })
+
+
+# ════════════════════════════════════════════════════════════════════════════════
+# Observation formatting (mirrors inference.py logic)
+# ════════════════════════════════════════════════════════════════════════════════
+
+MAX_QUEUE_NORM = 200.0
+MAX_LATENCY_NORM = 1000.0
+MAX_REQUEST_RATE_NORM = 100.0
+ALPHA, BETA, GAMMA, DELTA = 0.002, 0.01, 10.0, 0.005
+
+
+def format_observation(nodes: List[dict], task_id: str, step: int, max_steps: int) -> str:
+    """
+    Format simulator state as a text prompt for the model.
+
+    This mirrors inference.py's build_user_prompt and observation_for_model.
+    """
+    import json
+    node_data = []
+    for n in nodes:
+        node_data.append({
+            "node_id": n["node_id"],
+            "status": n["status"] if isinstance(n["status"], str) else n["status"].value,
+            "is_vip": n.get("is_vip", False),
+            "queue_depth": min(1.0, max(0.0, n["queue_depth"] / MAX_QUEUE_NORM)),
+            "latency_ms": min(1.0, max(0.0, n["latency_ms"] / MAX_LATENCY_NORM)),
+            "cpu_utilization": min(1.0, max(0.0, n.get("cpu_utilization", 0.0))),
+            "incoming_request_rate": min(1.0, max(0.0, n["incoming_request_rate"] / MAX_REQUEST_RATE_NORM)),
+        })
+    obs = {"task_id": task_id, "step": step, "max_steps": max_steps, "nodes": node_data}
+    return json.dumps(obs, separators=(",", ":"))
+
+
+def parse_action(action_text: str) -> dict:
+    """Parse model output into an action dict."""
+    import json
+    try:
+        data = json.loads(action_text)
+        return {
+            "action_type": str(data.get("action_type", "NO_OP")).upper(),
+            "target_node_id": str(data.get("target_node_id", "node-0")),
+            "parameter": float(data.get("parameter", 0.0)),
+        }
+    except (json.JSONDecodeError, ValueError):
+        return {"action_type": "NO_OP", "target_node_id": "node-0", "parameter": 0.0}
+
+
+# ════════════════════════════════════════════════════════════════════════════════
+# Episode Collection
+# ════════════════════════════════════════════════════════════════════════════════
+
+@dataclass
+class StepRecord:
+    """A single step in an episode."""
+    prompt: str           # Observation formatted as text
+    action_text: str      # Model output (JSON string)
+    log_prob: float       # log π(action | prompt)
+    reward: float         # Raw reward for this step
+    reward_normalized: float  # Normalized reward
+
+
+@dataclass
+class EpisodeRecord:
+    """A complete episode trajectory."""
+    task_id: str
+    steps: List[StepRecord] = field(default_factory=list)
+    total_reward: float = 0.0
+    total_reward_normalized: float = 0.0
+    avg_reward: float = 0.0
+    sla_violations: int = 0
+    final_lyapunov: float = 0.0
+
+
+class EpisodeCollector:
+    """
+    Collects episodes by running the simulator with a policy model.
+
+    This is the bridge between the simulator (physics) and the training
+    pipeline (loss computation). It produces EpisodeRecords that feed
+    directly into the loss functions.
+    """
+
+    def __init__(self, config: TrainingConfig):
+        self._config = config
+        self._sim = ClusterSimulator(n_nodes=config.n_nodes)
+
+    def collect_episode(
+        self,
+        model: PolicyModel,
+        task_id: str,
+        seed: Optional[int] = None,
+    ) -> EpisodeRecord:
+        """Run one episode and collect step-level data."""
+        cfg = self._config
+        self._sim.reset(task_id=task_id, seed=seed)
+
+        record = EpisodeRecord(task_id=task_id)
+        prev_lyapunov = 0.0
+
+        for step in range(1, cfg.max_steps + 1):
+            # 1. Get observation
+            nodes_true = self._sim.state(for_agent=False)
+            nodes_obs = self._sim.state(for_agent=True)
+            prompt = format_observation(nodes_obs, task_id, step, cfg.max_steps)
+
+            # 2. Get action from model
+            action_text = model.generate(prompt)
+            log_prob = model.get_log_prob(prompt, action_text)
+
+            # 3. Apply action
+            action = parse_action(action_text)
+            class _A:
+                pass
+            a = _A()
+            a.action_type = action["action_type"]
+            a.target_node_id = action["target_node_id"]
+            a.parameter = action["parameter"]
+            self._sim.apply_action(a)
+
+            # 4. Tick
+            self._sim.tick()
+
+            # 5. Compute reward (mirrors environment.py)
+            nodes_true = self._sim.state(for_agent=False)
+            current_lyapunov = compute_lyapunov(nodes_true)
+
+            # Importance-weighted average latency
+            w_lat = 0.0
+            w_sum = 0.0
+            for n in nodes_true:
+                w = n.get("importance_weight", 1.0)
+                lat = MAX_LATENCY_NORM if n["status"] == NodeStatus.FAILED else n["latency_ms"]
+                w_lat += w * lat
+                w_sum += w
+            avg_lat_norm = min(1.0, max(0.0, (w_lat / w_sum / MAX_LATENCY_NORM) if w_sum > 0 else 1.0))
+
+            # Error rate
+            total_in = sum(n.get("incoming_request_rate", 0) * n.get("importance_weight", 1.0) for n in nodes_true)
+            total_drop = sum(n.get("dropped_requests", 0) * n.get("importance_weight", 1.0) for n in nodes_true)
+            error_rate = min(1.0, total_drop / total_in) if total_in > 0 else 0.0
+
+            sla_step = smooth_sla_penalty(avg_lat_norm, error_rate)
+            if avg_lat_norm > 0.20 or error_rate > 0.05:
+                record.sla_violations += 1
+
+            # Cost
+            total_cap = 0
+            for n in nodes_true:
+                if n["status"] != NodeStatus.FAILED:
+                    total_cap += int(n.get("capacity_units", 0)) + int(n.get("pending_capacity_units", 0))
+            cost = total_cap * COST_PER_CAPACITY_UNIT_PER_HOUR
+
+            barrier = compute_barrier(nodes_true)
+            raw_reward = compute_reward(
+                prev_lyapunov, current_lyapunov, cost, sla_step,
+                ALPHA, BETA, GAMMA, barrier, DELTA,
+            )
+            norm_reward = normalize_reward(raw_reward)
+
+            record.steps.append(StepRecord(
+                prompt=prompt,
+                action_text=action_text,
+                log_prob=log_prob,
+                reward=raw_reward,
+                reward_normalized=norm_reward,
+            ))
+            record.total_reward += raw_reward
+            record.total_reward_normalized += norm_reward
+            prev_lyapunov = current_lyapunov
+
+        record.avg_reward = record.total_reward / max(1, len(record.steps))
+        record.final_lyapunov = prev_lyapunov
+        return record
+
+    def collect_group(
+        self,
+        model: PolicyModel,
+        task_id: str,
+        k: int,
+        seed: Optional[int] = None,
+    ) -> List[EpisodeRecord]:
+        """
+        Collect K episodes from the same initial state (for GRPO/RLOO).
+
+        Uses the same seed for all K episodes so they start from the same
+        domain randomization, but different model samples produce different
+        trajectories.
+        """
+        return [self.collect_episode(model, task_id, seed=seed) for _ in range(k)]
+
+
+# ════════════════════════════════════════════════════════════════════════════════
+# Trainer
+# ════════════════════════════════════════════════════════════════════════════════
+
+class SRETrainer:
+    """
+    Main training orchestrator for AntiAtropos SRE agents.
+
+    Usage (local validation with MockPolicyModel):
+        config = TrainingConfig(loss_fn="reinforce_baseline")
+        trainer = SRETrainer(config)
+        model = MockPolicyModel()
+        metrics = trainer.train_step(model, task_id="task-1", seed=42)
+
+    Usage (Colab with real model):
+        config = TrainingConfig(loss_fn="grpo", n_samples_per_state=4)
+        trainer = SRETrainer(config)
+        model = QLoRAModel(...)  # Your transformers model
+        for epoch in range(num_epochs):
+            for task in config.tasks:
+                metrics = trainer.train_step(model, task_id=task)
+                model.update(metrics["loss"])  # Backprop
+    """
+
+    def __init__(self, config: TrainingConfig):
+        assert config.loss_fn in VALID_LOSSES, f"Unknown loss: {config.loss_fn}"
+        self._config = config
+        self._collector = EpisodeCollector(config)
+        self._running_reward_mean = 0.0
+        self._running_reward_var = 1.0
+        self._episode_count = 0
+
+    def train_step(
+        self,
+        model: PolicyModel,
+        task_id: str,
+        seed: Optional[int] = None,
+    ) -> dict:
+        """
+        Execute one training step: collect episode(s) → compute loss.
+
+        Returns a metrics dict with:
+            - loss: The computed loss value
+            - avg_reward: Average raw reward across the episode
+            - avg_norm_reward: Average normalized reward
+            - episode_length: Number of steps
+            - sla_violations: Number of SLA violations
+            - final_lyapunov: Lyapunov energy at episode end
+            - reward_mean/var: Running reward statistics
+        """
+        cfg = self._config
+
+        if cfg.loss_fn in (LOSS_GRPO, LOSS_RLOO):
+            return self._train_step_grouped(model, task_id, seed)
+        else:
+            return self._train_step_single(model, task_id, seed)
+
+    def _train_step_single(
+        self,
+        model: PolicyModel,
+        task_id: str,
+        seed: Optional[int] = None,
+    ) -> dict:
+        """Train step for REINFORCE / REINFORCE+baseline."""
+        cfg = self._config
+
+        # 1. Collect episode
+        episode = self._collector.collect_episode(model, task_id, seed=seed)
+
+        # 2. Extract rewards and log probs
+        rewards = [s.reward for s in episode.steps]
+        log_probs = [s.log_prob for s in episode.steps]
+
+        # 3. Update running reward stats
+        ep_mean, ep_var = compute_reward_stats(rewards)
+        self._running_reward_mean = (
+            (1 - cfg.reward_ema_alpha) * self._running_reward_mean
+            + cfg.reward_ema_alpha * ep_mean
+        )
+        self._running_reward_var = (
+            (1 - cfg.reward_ema_alpha) * self._running_reward_var
+            + cfg.reward_ema_alpha * ep_var
+        )
+
+        # 4. Optionally normalize rewards
+        if cfg.normalize_rewards:
+            rewards = normalize_rewards(
+                rewards, self._running_reward_mean, self._running_reward_var
+            )
+
+        # 5. Compute returns
+        returns = compute_returns(rewards, gamma=cfg.gamma)
+
+        # 6. Compute loss
+        if cfg.loss_fn == LOSS_REINFORCE:
+            loss = reinforce_loss(log_probs, returns)
+        elif cfg.loss_fn == LOSS_REINFORCE_BASELINE:
+            # Use running mean as baseline
+            baselines = [self._running_reward_mean] * len(returns)
+            loss = reinforce_baseline_loss(
+                log_probs, returns, baselines,
+                normalize_advantage=cfg.normalize_advantages,
+            )
+        else:
+            raise ValueError(f"Unexpected loss_fn: {cfg.loss_fn}")
+
+        self._episode_count += 1
+
+        return {
+            "loss": loss,
+            "avg_reward": episode.avg_reward,
+            "avg_norm_reward": episode.total_reward_normalized / max(1, len(episode.steps)),
+            "episode_length": len(episode.steps),
+            "sla_violations": episode.sla_violations,
+            "final_lyapunov": episode.final_lyapunov,
+            "reward_mean": self._running_reward_mean,
+            "reward_var": self._running_reward_var,
+            "task_id": task_id,
+            "episode": episode,
+        }
+
+    def _train_step_grouped(
+        self,
+        model: PolicyModel,
+        task_id: str,
+        seed: Optional[int] = None,
+    ) -> dict:
+        """Train step for GRPO / RLOO."""
+        cfg = self._config
+        k = cfg.n_samples_per_state
+
+        # 1. Collect K episodes (same seed → same domain randomization)
+        episodes = self._collector.collect_group(model, task_id, k=k, seed=seed)
+
+        # 2. For each step position, form groups across episodes
+        #    (assumes all episodes have same length)
+        min_len = min(len(ep.steps) for ep in episodes)
+
+        log_probs_groups = []
+        rewards_groups = []
+
+        for t in range(min_len):
+            step_lps = []
+            step_rs = []
+            for ep in episodes:
+                step_lps.append(ep.steps[t].log_prob)
+                step_rs.append(ep.steps[t].reward)
+            log_probs_groups.append(step_lps)
+            rewards_groups.append(step_rs)
+
+        # 3. Update running stats
+        all_rewards = [s.reward for ep in episodes for s in ep.steps]
+        ep_mean, ep_var = compute_reward_stats(all_rewards)
+        self._running_reward_mean = (
+            (1 - cfg.reward_ema_alpha) * self._running_reward_mean
+            + cfg.reward_ema_alpha * ep_mean
+        )
+        self._running_reward_var = (
+            (1 - cfg.reward_ema_alpha) * self._running_reward_var
+            + cfg.reward_ema_alpha * ep_var
+        )
+
+        # 4. Normalize rewards
+        if cfg.normalize_rewards:
+            rewards_groups = [
+                normalize_rewards(rs, self._running_reward_mean, self._running_reward_var)
+                for rs in rewards_groups
+            ]
+
+        # 5. Compute loss
+        if cfg.loss_fn == LOSS_GRPO:
+            loss = grpo_loss(log_probs_groups, rewards_groups)
+        elif cfg.loss_fn == LOSS_RLOO:
+            loss = rloo_loss(log_probs_groups, rewards_groups)
+        else:
+            raise ValueError(f"Unexpected grouped loss_fn: {cfg.loss_fn}")
+
+        # 6. Aggregate metrics across episodes
+        avg_reward = sum(ep.avg_reward for ep in episodes) / len(episodes)
+        avg_norm = sum(
+            ep.total_reward_normalized / max(1, len(ep.steps)) for ep in episodes
+        ) / len(episodes)
+        total_sla = sum(ep.sla_violations for ep in episodes)
+        avg_lyapunov = sum(ep.final_lyapunov for ep in episodes) / len(episodes)
+
+        self._episode_count += k
+
+        return {
+            "loss": loss,
+            "avg_reward": avg_reward,
+            "avg_norm_reward": avg_norm,
+            "episode_length": min_len,
+            "sla_violations": total_sla,
+            "final_lyapunov": avg_lyapunov,
+            "reward_mean": self._running_reward_mean,
+            "reward_var": self._running_reward_var,
+            "task_id": task_id,
+            "episodes": episodes,
+        }
+
+    def train_epoch(
+        self,
+        model: PolicyModel,
+        seed: Optional[int] = None,
+    ) -> List[dict]:
+        """
+        Run one training step per task in the curriculum.
+
+        Returns a list of metrics dicts (one per task).
+        """
+        results = []
+        for task_id in self._config.tasks:
+            step_seed = seed + hash(task_id) % 1000 if seed is not None else None
+            metrics = self.train_step(model, task_id, seed=step_seed)
+            results.append(metrics)
+            if self._episode_count % self._config.log_every == 0:
+                self._log_metrics(metrics)
+        return results
+
+    def _log_metrics(self, metrics: dict) -> None:
+        """Print training metrics."""
+        print(
+            f"[Episode {self._episode_count}] "
+            f"task={metrics['task_id']} "
+            f"loss={metrics['loss']:.4f} "
+            f"avg_reward={metrics['avg_reward']:.4f} "
+            f"avg_norm_reward={metrics['avg_norm_reward']:.4f} "
+            f"sla_violations={metrics['sla_violations']} "
+            f"lyapunov={metrics['final_lyapunov']:.1f} "
+            f"reward_mean={metrics['reward_mean']:.4f} "
+            f"reward_var={metrics['reward_var']:.4f}"
+        )