[Contrib] Agent-OS Integration: Kernel-Level Safety for RL Training (#478)

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: imran-siddique

contrib/agentlightning/contrib/adapter/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+# Namespace package for agentlightning.contrib.adapter.

contrib/agentlightning/contrib/adapter/agentos.py

@@ -0,0 +1,127 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""
+FlightRecorderAdapter - Import Audit Logs to LightningStore
+=============================================================
+
+Adapts Agent-OS Flight Recorder to Agent-Lightning store format.
+"""
+
+from __future__ import annotations
+
+import logging
+from datetime import datetime, timezone
+from typing import Any, Dict, List
+
+logger = logging.getLogger(__name__)
+
+
+class FlightRecorderAdapter:
+    """
+    Import Agent-OS Flight Recorder logs to LightningStore.
+
+    Example:
+        >>> from agent_os import FlightRecorder
+        >>>
+        >>> recorder = FlightRecorder()
+        >>> adapter = FlightRecorderAdapter(recorder)
+        >>>
+        >>> # Import to Lightning store
+        >>> adapter.import_to_store(lightning_store)
+    """
+
+    def __init__(
+        self,
+        flight_recorder: Any,
+        *,
+        trace_id_prefix: str = "agentos",
+    ):
+        """
+        Initialize adapter.
+
+        Args:
+            flight_recorder: Agent-OS FlightRecorder
+            trace_id_prefix: Prefix for trace IDs
+        """
+        self.recorder = flight_recorder
+        self.trace_id_prefix = trace_id_prefix
+        self._imported_count = 0
+
+    def _convert_entry(self, entry: Any, index: int) -> Dict[str, Any]:
+        """Convert Flight Recorder entry to span format."""
+        entry_type = getattr(entry, "type", "unknown")
+        timestamp = getattr(entry, "timestamp", datetime.now(timezone.utc))
+        agent_id = getattr(entry, "agent_id", "unknown")
+
+        span = {
+            "span_id": f"{self.trace_id_prefix}-{index}",
+            "trace_id": f"{self.trace_id_prefix}-{agent_id}",
+            "name": f"agent_os.{entry_type}",
+            "start_time": timestamp.isoformat() if hasattr(timestamp, "isoformat") else str(timestamp),
+            "attributes": {
+                "agent_os.entry_type": entry_type,
+                "agent_os.agent_id": agent_id,
+            },
+        }
+
+        # Add type-specific attributes
+        if entry_type == "policy_check":
+            span["attributes"].update(
+                {
+                    "agent_os.policy_name": getattr(entry, "policy_name", "unknown"),
+                    "agent_os.policy_violated": getattr(entry, "violated", False),
+                }
+            )
+        elif entry_type == "signal":
+            span["attributes"].update(
+                {
+                    "agent_os.signal_type": getattr(entry, "signal", "unknown"),
+                }
+            )
+
+        return span
+
+    def get_spans(self) -> List[Dict[str, Any]]:
+        """Get all entries as spans."""
+        entries = []
+        if hasattr(self.recorder, "get_entries"):
+            entries = self.recorder.get_entries()
+        elif hasattr(self.recorder, "entries"):
+            entries = self.recorder.entries
+
+        return [self._convert_entry(e, i) for i, e in enumerate(entries)]
+
+    def import_to_store(self, store: Any) -> int:
+        """
+        Import spans to LightningStore.
+
+        Args:
+            store: LightningStore instance
+
+        Returns:
+            Number of spans imported
+        """
+        spans = self.get_spans()
+
+        for span in spans:
+            try:
+                if hasattr(store, "emit_span"):
+                    store.emit_span(span)
+                elif hasattr(store, "add_span"):
+                    store.add_span(span)
+            except Exception as e:
+                logger.error(f"Failed to import span: {e}")
+
+        self._imported_count += len(spans)
+        logger.info(f"Imported {len(spans)} spans to LightningStore")
+        return len(spans)
+
+    def get_violation_summary(self) -> Dict[str, Any]:
+        """Get summary of policy violations."""
+        spans = self.get_spans()
+        violations = [s for s in spans if s["attributes"].get("agent_os.policy_violated", False)]
+        return {
+            "total_entries": len(spans),
+            "total_violations": len(violations),
+            "violation_rate": len(violations) / len(spans) if len(spans) > 0 else 0.0,
+        }

contrib/agentlightning/contrib/reward/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+# Namespace package for agentlightning.contrib.reward.

contrib/agentlightning/contrib/reward/agentos.py

@@ -0,0 +1,127 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""
+PolicyReward - Convert Policy Violations to RL Penalties
+=========================================================
+
+Reward function that integrates Agent-OS governance.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Callable, Dict, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class PolicyReward:
+    """
+    Reward function that penalizes policy violations.
+
+    Example:
+        >>> from agent_os import KernelSpace
+        >>>
+        >>> kernel = KernelSpace(policy="strict")
+        >>> reward_fn = PolicyReward(kernel, base_reward_fn=accuracy)
+        >>>
+        >>> reward = reward_fn(rollout)  # Base reward - violation penalties
+    """
+
+    def __init__(
+        self,
+        kernel: Any,
+        *,
+        base_reward_fn: Optional[Callable[[Any], float]] = None,
+        critical_penalty: float = -100.0,
+        high_penalty: float = -50.0,
+        medium_penalty: float = -10.0,
+        low_penalty: float = -1.0,
+        clean_bonus: float = 5.0,
+    ):
+        """
+        Initialize policy-aware reward.
+
+        Args:
+            kernel: Agent-OS KernelSpace
+            base_reward_fn: Base reward function
+            critical_penalty: Penalty for critical violations
+            high_penalty: Penalty for high violations
+            medium_penalty: Penalty for medium violations
+            low_penalty: Penalty for low violations
+            clean_bonus: Bonus for clean execution
+        """
+        self.kernel = kernel
+        self.base_reward_fn = base_reward_fn or self._default_reward
+        self.penalties = {
+            "critical": critical_penalty,
+            "high": high_penalty,
+            "medium": medium_penalty,
+            "low": low_penalty,
+        }
+        self.clean_bonus = clean_bonus
+
+        self._total_rewards = 0
+        self._total_penalties = 0.0
+
+    def _default_reward(self, rollout: Any) -> float:
+        """Default: 1.0 for success, 0.0 for failure."""
+        return 1.0 if getattr(rollout, "success", False) else 0.0
+
+    def __call__(self, rollout: Any, *, emit: bool = True) -> float:
+        """
+        Calculate reward with policy penalties.
+
+        Args:
+            rollout: Rollout with violations attribute
+            emit: Emit reward span
+
+        Returns:
+            Final reward
+        """
+        base = self.base_reward_fn(rollout)
+
+        violations = getattr(rollout, "violations", [])
+        penalty = sum(self.penalties.get(v.severity, -10.0) for v in violations)
+
+        reward = base + penalty
+        if not violations:
+            reward += self.clean_bonus
+
+        self._total_rewards += 1
+        self._total_penalties += penalty
+
+        if emit:
+            self._emit_reward(reward, base, penalty, len(violations))
+
+        return reward
+
+    def _emit_reward(
+        self,
+        final: float,
+        base: float,
+        penalty: float,
+        violation_count: int,
+    ) -> None:
+        """Emit multi-dimensional reward."""
+        try:
+            from agentlightning.emitter import emit_reward
+
+            emit_reward(
+                {"final": final, "base": base, "policy_penalty": penalty},
+                primary_key="final",
+                attributes={"agent_os.violations": violation_count},
+            )
+        except ImportError:
+            logger.debug(
+                "agentlightning.emitter not available; skipping reward emission.",
+                exc_info=True,
+            )
+
+    def get_stats(self) -> Dict[str, float]:
+        """Get reward statistics."""
+        total = self._total_rewards or 1
+        return {
+            "total_rewards": self._total_rewards,
+            "avg_penalty": self._total_penalties / total,
+        }

contrib/agentlightning/contrib/runner/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+# Namespace package for agentlightning.contrib.runner.