Add reward handling and normalization features to environment

- Introduced raw and normalized reward attributes in ClusterObservation and AntiAtroposEnvironment.
- Implemented reward normalization logic in stability module.
- Updated observability metrics to track raw and normalized rewards.
- Added tests for reward normalization and environment reward behavior.

client.py

@@ -111,6 +111,9 @@ class AntiAtroposEnv(
             data_freshness_ms=obs_data.get("data_freshness_ms", 0),
             action_ack_status=obs_data.get("action_ack_status", "success"),
             choke_level=obs_data.get("choke_level", 0.0),
+            raw_reward=obs_data.get("raw_reward", 0.0),
+            normalized_reward=obs_data.get("normalized_reward", 0.0),
+            reward_scale_version=obs_data.get("reward_scale_version", "sigmoid-v1"),
             done=payload.get("done", False),
             reward=payload.get("reward", 0.0),
         )

models.py

@@ -155,6 +155,9 @@ class ClusterObservation(BaseModel):
     action_id: str = ""
     executor_latency_ms: float = Field(default=0.0, ge=0.0)
     executor_error_code: str = ""
+    raw_reward: float = 0.0
+    normalized_reward: float = Field(default=0.0, ge=0.0, le=1.0)
+    reward_scale_version: str = "sigmoid-v1"
     choke_level: float = 0.0
 
     nodes: list[NodeObservation]

server/AntiAtropos_environment.py

@@ -1,5 +1,6 @@
 import time
 import json
+import os
 import logging
 from uuid import uuid4
 
@@ -9,13 +10,13 @@ from openenv.core.env_server.types import State
 try:
     from ..models import SREAction, ClusterObservation, NodeObservation, NodeStatus, EnvironmentMode
     from ..simulator import ClusterSimulator, COST_PER_CAPACITY_UNIT_PER_HOUR
-    from ..stability import compute_lyapunov, compute_reward
+    from ..stability import compute_lyapunov, compute_reward, normalize_reward, REWARD_SCALE_VERSION
     from ..telemetry import PrometheusClient, get_observability_tracker
     from ..control import KubernetesExecutor, ActionValidator
 except ImportError:
     from models import SREAction, ClusterObservation, NodeObservation, NodeStatus, EnvironmentMode  # type: ignore[no-redef]
     from simulator import ClusterSimulator, COST_PER_CAPACITY_UNIT_PER_HOUR  # type: ignore[no-redef]
-    from stability import compute_lyapunov, compute_reward  # type: ignore[no-redef]
+    from stability import compute_lyapunov, compute_reward, normalize_reward, REWARD_SCALE_VERSION  # type: ignore[no-redef]
     from telemetry import PrometheusClient, get_observability_tracker  # type: ignore[no-redef]
     from control import KubernetesExecutor, ActionValidator  # type: ignore[no-redef]
 
@@ -34,6 +35,7 @@ MAX_REQUEST_RATE_NORM = 100.0
 
 MAX_STEPS: int = 100      # Episode length
 N_NODES:   int = 5        # Cluster size
+REWARD_OUTPUT_MODES = {"normalized", "raw"}
 
 
 class AntiAtroposEnvironment(Environment):
@@ -70,6 +72,11 @@ class AntiAtroposEnvironment(Environment):
         self._last_action_id: str = ""
         self._last_executor_latency_ms: float = 0.0
         self._last_executor_error_code: str = ""
+        self._last_raw_reward: float = 0.0
+        self._last_normalized_reward: float = 0.0
+        self._reward_output_mode: str = os.getenv("ANTIATROPOS_REWARD_OUTPUT_MODE", "normalized").strip().lower()
+        if self._reward_output_mode not in REWARD_OUTPUT_MODES:
+            self._reward_output_mode = "normalized"
         self._last_metric_time: float = 0.0
 
     def reset(self, task_id: str = "task-1", mode: str = "simulated") -> ClusterObservation:
@@ -88,6 +95,8 @@ class AntiAtroposEnvironment(Environment):
         self._last_action_id = ""
         self._last_executor_latency_ms = 0.0
         self._last_executor_error_code = ""
+        self._last_raw_reward = 0.0
+        self._last_normalized_reward = 0.0
         
         # Only set baseline metric time for hybrid/live to prevent misleading freshness in SIM
         if self._mode in [EnvironmentMode.HYBRID, EnvironmentMode.LIVE]:
@@ -123,6 +132,8 @@ class AntiAtroposEnvironment(Environment):
         self._last_action_id = str(uuid4())
         self._last_executor_latency_ms = 0.0
         self._last_executor_error_code = ""
+        self._last_raw_reward = 0.0
+        self._last_normalized_reward = 0.0
         
         # 1. Action Validation & Execution
         valid_targets = [n["node_id"] for n in self._nodes_true]
@@ -196,7 +207,7 @@ class AntiAtroposEnvironment(Environment):
         
         # 7. Compute scalar reward
         cost = self._compute_cost(self._nodes_true)
-        reward = compute_reward(
+        raw_reward = compute_reward(
             v_prev=self._prev_lyapunov,
             v_curr=current_lyapunov,
             cost=cost,
@@ -205,6 +216,10 @@ class AntiAtroposEnvironment(Environment):
             beta=BETA,
             gamma=GAMMA
         )
+        normalized_reward = normalize_reward(raw_reward)
+        reward = normalized_reward if self._reward_output_mode == "normalized" else raw_reward
+        self._last_raw_reward = raw_reward
+        self._last_normalized_reward = normalized_reward
         
         self._prev_lyapunov = current_lyapunov
 
@@ -225,7 +240,9 @@ class AntiAtroposEnvironment(Environment):
             action_type=str(action.action_type.value),
             target_node_id=str(action.target_node_id),
             ack_status=self._action_ack_status,
-            reward=reward,
+            reward_output=reward,
+            reward_raw=raw_reward,
+            reward_normalized=normalized_reward,
             lyapunov_energy=obs.lyapunov_energy,
             total_queue_backlog=obs.total_queue_backlog,
             average_latency_ms=obs.average_latency_ms,
@@ -248,7 +265,10 @@ class AntiAtroposEnvironment(Environment):
                     "action_ack_status": self._action_ack_status,
                     "executor_latency_ms": self._last_executor_latency_ms,
                     "executor_error_code": self._last_executor_error_code,
-                    "reward": reward,
+                    "reward_output": reward,
+                    "reward_raw": raw_reward,
+                    "reward_normalized": normalized_reward,
+                    "reward_output_mode": self._reward_output_mode,
                     "lyapunov_energy": obs.lyapunov_energy,
                     "average_latency_ms_norm": obs.average_latency_ms,
                     "total_queue_backlog_norm": obs.total_queue_backlog,
@@ -364,6 +384,9 @@ class AntiAtroposEnvironment(Environment):
             action_id=self._last_action_id,
             executor_latency_ms=self._last_executor_latency_ms,
             executor_error_code=self._last_executor_error_code,
+            raw_reward=self._last_raw_reward,
+            normalized_reward=self._last_normalized_reward,
+            reward_scale_version=REWARD_SCALE_VERSION,
             choke_level=0.0,
             done=False,
             reward=0.0,

stability.py

@@ -38,6 +38,7 @@ Key concepts implemented
 
 from __future__ import annotations
 
+import os
 import math
 import statistics
 from typing import Sequence
@@ -57,6 +58,16 @@ STABILITY_WINDOW: int = 10
 trend-stable (V is on a decreasing trajectory)."""
 
 
+# ---------------------------------------------------------------------------
+# Reward normalisation defaults (env-overridable)
+# ---------------------------------------------------------------------------
+
+REWARD_NORM_MIDPOINT: float = float(os.getenv("ANTIATROPOS_REWARD_MIDPOINT", "0.0"))
+REWARD_NORM_TEMPERATURE: float = float(os.getenv("ANTIATROPOS_REWARD_TEMPERATURE", "5.0"))
+REWARD_NORM_EPS: float = float(os.getenv("ANTIATROPOS_REWARD_EPS", "1e-8"))
+REWARD_SCALE_VERSION: str = "sigmoid-v1"
+
+
 # ---------------------------------------------------------------------------
 # Core Lyapunov functions
 # ---------------------------------------------------------------------------
@@ -264,3 +275,25 @@ def compute_reward(
     """
     delta_v = compute_drift(v_prev, v_curr)
     return -(alpha * delta_v + beta * cost + gamma * sla_violation_step)
+
+
+def normalize_reward(
+    raw_reward: float,
+    midpoint: float = REWARD_NORM_MIDPOINT,
+    temperature: float = REWARD_NORM_TEMPERATURE,
+    eps: float = REWARD_NORM_EPS,
+) -> float:
+    """
+    Deterministically map raw reward to [0, 1] using a smooth sigmoid.
+
+    reward_01 = 1 / (1 + exp(-(raw_reward - midpoint) / temperature))
+    """
+    temp = max(float(eps), abs(float(temperature)))
+    z = (float(raw_reward) - float(midpoint)) / temp
+    if z >= 0:
+        exp_neg = math.exp(-z)
+        out = 1.0 / (1.0 + exp_neg)
+    else:
+        exp_pos = math.exp(z)
+        out = exp_pos / (1.0 + exp_pos)
+    return min(1.0, max(0.0, float(out)))

telemetry/observability.py

@@ -47,7 +47,17 @@ class ObservabilityTracker:
 
         self.reward_gauge = Gauge(
             "antiatropos_reward",
-            "Latest reward value",
+            "Latest output reward value (depends on reward output mode)",
+            ["task_id", "mode"],
+        )
+        self.reward_raw_gauge = Gauge(
+            "antiatropos_reward_raw",
+            "Latest raw reward value before normalization",
+            ["task_id", "mode"],
+        )
+        self.reward_normalized_gauge = Gauge(
+            "antiatropos_reward_normalized",
+            "Latest normalized reward value in [0,1]",
             ["task_id", "mode"],
         )
         self.lyapunov_gauge = Gauge(
@@ -73,7 +83,9 @@ class ObservabilityTracker:
         action_type: str,
         target_node_id: str,
         ack_status: str,
-        reward: float,
+        reward_output: float,
+        reward_raw: float,
+        reward_normalized: float,
         lyapunov_energy: float,
         total_queue_backlog: float,
         average_latency_ms: float,
@@ -93,7 +105,9 @@ class ObservabilityTracker:
                 target_node_id=target_node_id,
                 ack_class=ack_class,
             ).inc()
-            self.reward_gauge.labels(task_id=task_id, mode=mode).set(float(reward))
+            self.reward_gauge.labels(task_id=task_id, mode=mode).set(float(reward_output))
+            self.reward_raw_gauge.labels(task_id=task_id, mode=mode).set(float(reward_raw))
+            self.reward_normalized_gauge.labels(task_id=task_id, mode=mode).set(float(reward_normalized))
             self.lyapunov_gauge.labels(task_id=task_id, mode=mode).set(float(lyapunov_energy))
             self.queue_gauge.labels(task_id=task_id, mode=mode).set(float(total_queue_backlog))
             self.latency_gauge.labels(task_id=task_id, mode=mode).set(float(average_latency_ms))