Support for trajectory interfaces

verl/utils/tracking.py

@@ -264,9 +264,7 @@ def init(cls, project_name=None, experiment_name=None, config=None):
     def log(cls, data, step):
         if not cls.enabled():
             return
-        if not cls._init_done:
-            cls._get_rl_insight().init()
-            cls._init_done = True
+        cls._ensure_rl_insight_init()
         metric_gauge = cls._get_rl_insight().metric_gauge
 
         for key, value in data.items():
@@ -298,11 +296,69 @@ def trace_state(
             yield
             return
 
+        cls._ensure_rl_insight_init()
+        with cls._get_rl_insight().trace_state(state_name, state_lane_id=state_lane_id, **labels):
+            yield
+
+    @classmethod
+    def trace_span(
+        cls,
+        name: str,
+        *,
+        start_time_ns: int,
+        end_time_ns: int,
+        attributes: dict[str, Any] | None = None,
+    ) -> None:
+        """Report one completed span through rl-insight's direct trace API."""
+        if not cls.enabled():
+            return
+
+        cls._ensure_rl_insight_init()
+        cls._get_rl_insight().trace_span(
+            name=name,
+            start_time_ns=start_time_ns,
+            end_time_ns=end_time_ns,
+            attributes=dict(attributes or {}),
+        )
+
+    @classmethod
+    def _ensure_rl_insight_init(cls) -> None:
         if not cls._init_done:
             cls._get_rl_insight().init()
             cls._init_done = True
-        with cls._get_rl_insight().trace_state(state_name, state_lane_id=state_lane_id, **labels):
-            yield
+
+    @classmethod
+    def agent_loop_lane_id(cls, run_id: Any, sample: Any, session: Any, traj: Any) -> str:
+        from rl_insight.agent_loop import agent_loop_lane_id
+
+        return agent_loop_lane_id(run_id, sample, session, traj)
+
+    @classmethod
+    def publish_agent_loop_session(
+        cls,
+        *,
+        run_id: Any,
+        sample: Any,
+        session: Any,
+        trajectories: list[Any],
+        start_time_ns: int | None = None,
+        end_time_ns: int | None = None,
+    ) -> None:
+        """Publish dashboard hierarchy gauges for one finalized session."""
+        if not cls.enabled():
+            return
+
+        cls._ensure_rl_insight_init()
+        from rl_insight.agent_loop import publish_agent_loop_session
+
+        publish_agent_loop_session(
+            run_id=run_id,
+            sample=sample,
+            session=session,
+            trajectories=trajectories,
+            start_time_ns=start_time_ns,
+            end_time_ns=end_time_ns,
+        )
 
     @classmethod
     def register_rollout_metrics(