github-featured/microsoft--agent-lightning
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity

Bug fixes and RL improvements for Agent-framework support (#96)

Browse Source
This commit is contained in:
Yuge Zhang
2025-09-18 21:22:09 +08:00
committed by GitHub
parent 7ad967daf7
commit ba10c845e1
4 changed files with 58 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
agentlightning/runner.py
+3 -1
View File
@@ -243,7 +243,9 @@ class AgentRunner(ParallelWorkerBase):
end_time = time.time()
logger.info(
f"{self._log_prefix(rollout_id)} Completed in "
f"{end_time - start_time:.2f}s. Reward: {rollout_obj.final_reward}"
f"{end_time - start_time:.2f}s. Triplet length: "
f"{len(rollout_obj.triplets) if rollout_obj.triplets is not None else 'N/A'}. "
f"Reward: {rollout_obj.final_reward}"
)
except Exception:
logger.exception(f"{self._log_prefix(rollout_id)} Exception during rollout.")
agentlightning/tracer/triplet.py
+5
View File
@@ -259,6 +259,11 @@ class TraceTree:
if agent_name is not None:
return agent_name
# Case 5: agent-framework
agent_name = self.span.attributes.get("executor.id")
if agent_name is not None:
return agent_name
def maybe_reward_dict(self) -> dict[str, Any]:
for key in [
"agentops.task.output", # newer versions of agentops
agentlightning/verl/daemon.py
+48 -11
View File
@@ -5,6 +5,7 @@ import socket
import threading
import time
import uuid
from collections.abc import Mapping, Sequence
from typing import Dict, List, Optional
import numpy as np
@@ -84,6 +85,28 @@ def _find_available_port() -> int:
return s.getsockname()[1]
def _to_native(obj):
"""Convert data retrieved from Parquet to data usable in AGL server."""
# 1) Arrays -> list (then recurse)
if isinstance(obj, np.ndarray):
return _to_native(obj.tolist())
# 2) NumPy scalar types -> Python scalars
if isinstance(obj, np.generic):
return _to_native(obj.item())
# 3) Dict-like -> dict
if isinstance(obj, Mapping):
return {_to_native(k): _to_native(v) for k, v in obj.items()}
# 4) Lists/Tuples/Sets -> list
if isinstance(obj, (list, tuple, set)):
return [_to_native(x) for x in obj]
# 5) Anything else: leave as-is
return obj
class AgentModeDaemon:
"""
AgentModeDaemon using the AgentLightningServer SDK.
@@ -102,7 +125,7 @@ class AgentModeDaemon:
mini_batch_size,
pad_token_id,
reward_fillna_value=0.0,
llm_timeout_seconds=600.0,
llm_timeout_seconds=1200.0,
):
# Server and Task Configuration
self.server_port = port
@@ -270,7 +293,7 @@ class AgentModeDaemon:
# Data ID is different from Rollout ID, as one data can have multiple rollouts.
rollout_id = await self.server.queue_task(
sample=original_sample,
sample=_to_native(original_sample),
mode="train" if is_train else "val",
resources_id=resources_id,
metadata=task_metadata,
@@ -345,10 +368,12 @@ class AgentModeDaemon:
sample_stat_list = []
for rollout_id, rollout in self._completed_rollouts.items():
final_reward = self._fillna_reward(rollout)
if not rollout.triplets:
print(f"Warning: No triplets found for test rollout {rollout.rollout_id}.")
sample_stat_list.append({"reward": final_reward})
continue
response_length_list = [len(triplet.response.get("token_ids", [])) for triplet in rollout.triplets]
final_reward = self._fillna_reward(rollout)
sample_stat_list.append(
{
"sum_response_length": np.sum(response_length_list),
@@ -358,11 +383,16 @@ class AgentModeDaemon:
}
)
stats_w_trace = [stat for stat in sample_stat_list if "sum_response_length" in stat]
return {
"val/reward": np.mean([stat["reward"] for stat in sample_stat_list]),
"val/mean_response_length": np.mean([stat["mean_response_length"] for stat in sample_stat_list]),
"val/sum_response_length": np.mean([stat["sum_response_length"] for stat in sample_stat_list]),
"val/turn_count": np.mean([stat["turn_count"] for stat in sample_stat_list]),
"val/n_rollouts": len(sample_stat_list),
"val/n_rollouts_w_trace": len(stats_w_trace),
"val/reward": np.mean(
[stat["reward"] for stat in sample_stat_list]
), # each rollout must have a reward (fillna if missing)
"val/mean_response_length": np.mean([stat["mean_response_length"] for stat in stats_w_trace]),
"val/sum_response_length": np.mean([stat["sum_response_length"] for stat in stats_w_trace]),
"val/turn_count": np.mean([stat["turn_count"] for stat in stats_w_trace]),
}
def get_train_data_batch(self, max_prompt_length, max_response_length, device):
@@ -378,10 +408,15 @@ class AgentModeDaemon:
# 1. Reconstruct the `finished_id_to_sample_info` structure from completed rollouts
finished_id_to_sample_info = {}
finished_id_to_final_reward = {}
for rollout_id, rollout in self._completed_rollouts.items():
original_sample = self._task_id_to_original_sample[rollout_id]
final_reward = self._fillna_reward(rollout)
if not rollout.triplets:
finished_id_to_final_reward[rollout_id] = final_reward
print(f"Warning: No triplets found for training rollout {rollout.rollout_id}, skipping.")
continue
# The client should report triplets that contain prompt_ids and response_ids.
@@ -391,14 +426,13 @@ class AgentModeDaemon:
{"prompt_ids": t.prompt.get("token_ids", []), "response_ids": t.response.get("token_ids", [])}
for t in rollout.triplets
]
final_reward = self._fillna_reward(rollout)
info = {
"reward": final_reward,
"trace_list": trace_list,
"data_id": original_sample["data_id"],
}
finished_id_to_sample_info[rollout_id] = info
finished_id_to_final_reward[rollout_id] = final_reward
#
# --- Data processing and tensor creation logic ---
# Get all the reported data.
@@ -487,8 +521,11 @@ class AgentModeDaemon:
data_proto = DataProto(batch=batch)
data_metrics = {
"agent_mode/n_trunc_sample_because_of_response": n_trunc_sample_because_of_response,
"agent_mode/n_sample_to_train": n_transition,
"training/reward": np.mean(list(finished_id_to_final_reward.values())),
"training/n_rollouts": len(finished_id_to_final_reward),
"training/n_rollouts_w_trace": len(finished_id_to_sample_info),
"training/n_truncated_triplets": n_trunc_sample_because_of_response,
"training/n_triplets": n_transition,
}
# Add non-tensor data for advantage calculation and logging
agentlightning/verl/trainer.py
+2 -2
View File
@@ -195,7 +195,7 @@ class AgentLightningTrainer(RayPPOTrainer):
# after advantages are assinged, we begin to drop (1) long prompt (2) floor to ppo minisize
keep_indices = (~batch.batch["is_drop_mask"]).nonzero(as_tuple=True)[0]
metrics["agent_mode/n_dropped_sample_because_of_prompt"] = (
metrics["training/n_triplets_prompt_too_long"] = (
batch.batch["is_drop_mask"].shape[0] - keep_indices.shape[0]
)
batch = batch[keep_indices]
@@ -207,7 +207,7 @@ class AgentLightningTrainer(RayPPOTrainer):
batch.reorder(torch.tensor(random_indices).type(torch.int32))
n_remained_transition = n_transition // mini_batch_size * mini_batch_size
batch = batch[list(range(n_remained_transition))]
metrics["agent_mode/n_dropped_sample_because_of_mini_batch"] = n_transition - n_remained_transition
metrics["training/n_triplets_dropped_remainder"] = n_transition - n_remained_transition
# Agent mode note: Change the order of balance batch;
# 1. first calculate advantage