Skip to content

fix(train): analyze agent evolution once per commit without evaluator #6015

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
akushonkamen wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(train): analyze agent evolution once per commit without evaluator #6015

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
163 changes: 92 additions & 71 deletions openviking/session/compressor_v3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1038,66 +1038,104 @@ async def train_from_extracted_cases(

case_uri_map = dict(case_uri_by_name or {})

# Without a rollout evaluator the case never reaches the model: the
# analyzer only stamps case_name onto the extracted trajectories
# after extraction, so a per-case loop would rerun the identical
# whole-session extraction, gradient estimation and merge once per
# case. In that case run the pipeline once for the first case and
# link every case to the shared result. With an evaluator the
# per-case analysis is rubric-specific and the per-case loop is
# preserved.
share_training = getattr(self.rollout_analyzer, "evaluator", None) is None
shared_training: tuple[RolloutAnalysis, RolloutTrainingResult] | None = None

for case in cases:
case_uri = _case_uri_for_case(case, case_uri_map)
rollout = Rollout(
case=case,
messages=list(messages),
policy_snapshot_id=policy_snapshot_id,
)
# Analyze once — trajectories + skill patches co-extracted
analysis = await self.rollout_analyzer.analyze(rollout, analysis_context)
if share_training and shared_training is not None:
reused_shared = True
analysis, exp_training_result = shared_training
else:
reused_shared = False
rollout = Rollout(
case=case,
messages=list(messages),
policy_snapshot_id=policy_snapshot_id,
)
# Analyze once — trajectories + skill patches co-extracted
analysis = await self.rollout_analyzer.analyze(rollout, analysis_context)

# Experience path: estimate gradients, then submit to exp trainer
exp_gradients = []
if experiences_allowed:
exp_gradients = await ExperienceGradientEstimator(
viking_fs=viking_fs,
vlm_resolver=self.vlm_resolver,
).estimate(analysis, exp_trainer.policy_set, gradient_context)
exp_training_result = _trajectory_only_training_result(
analysis=analysis,
rollout=rollout,
policy_set=exp_trainer.policy_set,
)
if exp_gradients:
fallback_trajectory_uris = {
uri
for trajectory in analysis.trajectories
if (uri := str(getattr(trajectory, "uri", "") or ""))
}

# Experience path: estimate gradients, then submit to exp trainer
exp_gradients = []
if experiences_allowed:
exp_gradients = await ExperienceGradientEstimator(
viking_fs=viking_fs,
vlm_resolver=self.vlm_resolver,
).estimate(analysis, exp_trainer.policy_set, gradient_context)
exp_training_result = _trajectory_only_training_result(
analysis=analysis,
rollout=rollout,
policy_set=exp_trainer.policy_set,
)
if exp_gradients:
fallback_trajectory_uris = {
uri
for trajectory in analysis.trajectories
if (uri := str(getattr(trajectory, "uri", "") or ""))
}
async def commit_experience_batch(
batch_result: RolloutTrainingResult,
*,
_fallback_trajectory_uris: set[str] = fallback_trajectory_uris,
) -> None:
persisted_result = (
getattr(batch_result, "batch_result", None) or batch_result
)
snapshot_apply_result, experience_trajectory_map = (
_experience_snapshot_provenance(
batch_result,
fallback_trajectory_uris=_fallback_trajectory_uris,
)
)
await _commit_experience_snapshot(
viking_fs,
ctx=ctx,
experience_uris=_visible_experience_snapshot_uris(
plan=persisted_result.plan,
apply_result=snapshot_apply_result,
),
archive_uri=archive_uri,
experience_trajectory_map=experience_trajectory_map,
)

async def commit_experience_batch(
batch_result: RolloutTrainingResult,
*,
_fallback_trajectory_uris: set[str] = fallback_trajectory_uris,
) -> None:
persisted_result = (
getattr(batch_result, "batch_result", None) or batch_result
exp_training_result = await exp_trainer.submit_gradients(
exp_gradients,
analysis=analysis,
rollout=rollout,
batch_finalizer=commit_experience_batch,
)
snapshot_apply_result, experience_trajectory_map = (
_experience_snapshot_provenance(
batch_result,
fallback_trajectory_uris=_fallback_trajectory_uris,
if share_training:
shared_training = (analysis, exp_training_result)
# Skill path: co-extracted skill gradients go directly to skill
# trainer. On the shared path this runs once per commit with
# the single analysis; per-case analyses submit their own.
if skill_trainer is not None and analysis.gradients:
skill_gradients = [
g for g in analysis.gradients if _gradient_memory_type(g) == "skills"
]
if skill_gradients:
skill_training_result = await skill_trainer.submit_gradients(
skill_gradients,
analysis=analysis,
rollout=rollout,
)
)
await _commit_experience_snapshot(
viking_fs,
ctx=ctx,
experience_uris=_visible_experience_snapshot_uris(
plan=persisted_result.plan,
apply_result=snapshot_apply_result,
),
archive_uri=archive_uri,
experience_trajectory_map=experience_trajectory_map,
)
skill_submitted += 1
apply_result = getattr(skill_training_result, "apply_result", None)
if apply_result is not None:
for uri in getattr(apply_result, "written_uris", []) or []:
if uri:
skill_uris.append(str(uri))

exp_training_result = await exp_trainer.submit_gradients(
exp_gradients,
analysis=analysis,
rollout=rollout,
batch_finalizer=commit_experience_batch,
)
if case_uri:
await self._link_case_to_training_outputs(
analysis=analysis,
Expand All @@ -1107,27 +1145,10 @@ async def commit_experience_batch(
ctx=ctx,
viking_fs=viking_fs,
)
# Skill path: co-extracted skill gradients go directly to skill trainer
if skill_trainer is not None and analysis.gradients:
skill_gradients = [
g for g in analysis.gradients if _gradient_memory_type(g) == "skills"
]
if skill_gradients:
skill_training_result = await skill_trainer.submit_gradients(
skill_gradients,
analysis=analysis,
rollout=rollout,
)
skill_submitted += 1
apply_result = getattr(skill_training_result, "apply_result", None)
if apply_result is not None:
for uri in getattr(apply_result, "written_uris", []) or []:
if uri:
skill_uris.append(str(uri))

submitted += 1

if collect_memory_diff:
if collect_memory_diff and not reused_shared:
# Build diff from the strongly typed training result returned by
# submit_gradients. Do not use exp_trainer.last_apply_result here:
# it is only a PolicyApplyResult and does not carry analyses/plan,
Expand Down
177 changes: 170 additions & 7 deletions tests/session/test_compressor_v3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -784,16 +784,179 @@ async def fake_estimate_exp_gradients(self, *args, **kwargs):

assert result["case_count"] == 2
assert result["submitted"] == 2
assert [rollout.case.name for rollout in seen_rollouts] == [
"duplicate_booking",
"case_b",
]
assert [rollout.messages for rollout in seen_rollouts] == [
rollout_messages,
rollout_messages,
# Without an evaluator the whole-session analysis runs once for the commit;
# the analyzed rollout is bound to the first case and both cases are linked
# to the shared training result.
assert len(seen_rollouts) == 1
assert seen_rollouts[0].case.name == "duplicate_booking"
assert seen_rollouts[0].messages == rollout_messages


@pytest.mark.asyncio
async def test_train_from_extracted_multiple_cases_train_once_without_evaluator(monkeypatch):
"""Without an evaluator every case sees byte-identical model input, so the
whole-session analyze/estimate runs once per commit and each extracted case
is linked to the shared training result."""
analyze_calls: list[str] = []
estimate_calls: list[str] = []

class FakeTrainer:
policy_set = ExperienceSet(
root_uri="viking://user/u/memories/experiences",
policies=[],
)

class FakeAnalyzer:
evaluator = None # matches the analyzer SessionCompressorV3 builds by default

async def analyze(self, rollout, context):
del context
analyze_calls.append(rollout.case.name)
return RolloutAnalysis(
evaluation=RubricEvaluation(
passed=True,
score=1.0,
criterion_results=[],
feedback=[],
),
trajectories=[],
gradients=[],
)

async def fake_estimate_exp_gradients(self, *args, **kwargs):
estimate_calls.append("estimate")
return []

monkeypatch.setattr(
"openviking.session.compressor_v3.get_viking_fs",
lambda: SimpleNamespace(ls=AsyncMock(return_value=[])),
)
monkeypatch.setattr(
"openviking.session.compressor_v3.get_streaming_policy_trainer",
AsyncMock(return_value=FakeTrainer()),
)
monkeypatch.setattr(
"openviking.session.train.components.gradient_estimator.ExperienceGradientEstimator.estimate",
fake_estimate_exp_gradients,
)

case_a = _training_case()
case_b = Case(
name="case_b",
task_signature="Handle a second extracted case.",
input={"summary": "second case"},
rubric=case_a.rubric,
)

compressor = SessionCompressorV3(
vikingdb=None,
rollout_analyzer=FakeAnalyzer(),
streaming_trainer_config=StreamingPolicyTrainerConfig(
max_wait_seconds=60,
max_gradients_per_update=8,
),
)
link_mock = AsyncMock()
compressor._link_case_to_training_outputs = link_mock

result = await compressor.train_from_extracted_cases(
cases=[case_a, case_b],
messages=_messages(),
ctx=_ctx(),
session_id="s1",
case_uri_by_name={
"duplicate_booking": "viking://user/u/memories/cases/case_a.md",
"case_b": "viking://user/u/memories/cases/case_b.md",
},
)

assert result["case_count"] == 2
assert result["submitted"] == 2
# One whole-session analysis per commit, not one per case.
assert analyze_calls == ["duplicate_booking"]
assert len(estimate_calls) == 1
# Every extracted case is linked to the shared training result.
assert link_mock.await_count == 2
assert [call.kwargs["case_uri"] for call in link_mock.await_args_list] == [
"viking://user/u/memories/cases/case_a.md",
"viking://user/u/memories/cases/case_b.md",
]


@pytest.mark.asyncio
async def test_train_from_extracted_multiple_cases_keeps_per_case_loop_with_evaluator(monkeypatch):
"""With a rollout evaluator the per-case analysis is rubric-specific, so the
per-case loop is preserved."""
analyze_calls: list[str] = []

class FakeTrainer:
policy_set = ExperienceSet(
root_uri="viking://user/u/memories/experiences",
policies=[],
)

class FakeAnalyzer:
evaluator = SimpleNamespace(evaluate=AsyncMock(return_value=None))

async def analyze(self, rollout, context):
del context
analyze_calls.append(rollout.case.name)
return RolloutAnalysis(
evaluation=RubricEvaluation(
passed=True,
score=1.0,
criterion_results=[],
feedback=[],
),
trajectories=[],
gradients=[],
)

async def fake_estimate_exp_gradients(self, *args, **kwargs):
return []

monkeypatch.setattr(
"openviking.session.compressor_v3.get_viking_fs",
lambda: SimpleNamespace(ls=AsyncMock(return_value=[])),
)
monkeypatch.setattr(
"openviking.session.compressor_v3.get_streaming_policy_trainer",
AsyncMock(return_value=FakeTrainer()),
)
monkeypatch.setattr(
"openviking.session.train.components.gradient_estimator.ExperienceGradientEstimator.estimate",
fake_estimate_exp_gradients,
)

case_a = _training_case()
case_b = Case(
name="case_b",
task_signature="Handle a second extracted case.",
input={"summary": "second case"},
rubric=case_a.rubric,
)

compressor = SessionCompressorV3(
vikingdb=None,
rollout_analyzer=FakeAnalyzer(),
streaming_trainer_config=StreamingPolicyTrainerConfig(
max_wait_seconds=60,
max_gradients_per_update=8,
),
)

result = await compressor.train_from_extracted_cases(
cases=[case_a, case_b],
messages=_messages(),
ctx=_ctx(),
session_id="s1",
)

assert result["case_count"] == 2
assert result["submitted"] == 2
assert analyze_calls == ["duplicate_booking", "case_b"]


@pytest.mark.asyncio
async def test_v3_extract_uses_patch_merge_without_directory_lock(monkeypatch):
applied_operations = []
Expand Down
Loading