From 6935003ff11015bfd937ade361066c14fda12e02 Mon Sep 17 00:00:00 2001 From: akushonkamen Date: Tue, 13 Oct 2026 20:15:09 +0800 Subject: [PATCH] fix(train): analyze agent evolution once per commit without evaluator Without a rollout evaluator the extracted case never reaches the model: TrajectoryRolloutAnalyzer only stamps case_name onto the extracted trajectories after extraction, so training once per case reran the identical whole-session extraction, gradient estimation and patch merge N times per commit (thinking-enabled model calls included). Run the pipeline once for the first case and link every case to the shared training result; with an evaluator the per-case rubric-specific loop is preserved. Test evidence (local minimal OPENVIKING_CONFIG_FILE; all pytest runs with -q --no-cov): - red = old source (938ef31) + this commit's test file: tests/session/test_compressor_v3.py -> 8 failed / 28 passed. Beyond the 6 pre-existing environment failures, both test_train_from_extracted_multiple_cases_train_once_without_evaluator (analyze ran twice for 2 cases) and the updated test_train_from_extracted_multiple_case_memories_analyzes_bound_rollouts (asserts one whole-session analysis) fail on the old code. - green = this commit: same file -> 6 failed / 30 passed; the 6 failures are identical to the pristine 938ef31 baseline (pre-existing, e.g. stale FakeTrainer without the batch_finalizer kwarg). - adjacency: tests/session/train/ -> 23 failed / 87 passed, failure list identical to pristine baseline; test_trajectory_analyzer_component.py + test_gradient_estimator_component.py -> 8 passed; test_session_commit.py -> 20 passed. - ruff check clean on both touched files; mypy error list unchanged (2 pre-existing errors shifted line numbers only). Refs #6006 Signed-off-by: akushonkamen --- openviking/session/compressor_v3.py | 163 ++++++++++++++----------- tests/session/test_compressor_v3.py | 177 ++++++++++++++++++++++++++-- 2 files changed, 262 insertions(+), 78 deletions(-) diff --git a/openviking/session/compressor_v3.py b/openviking/session/compressor_v3.py index 0a4f2c5b66..a288729964 100644 --- a/openviking/session/compressor_v3.py +++ b/openviking/session/compressor_v3.py @@ -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, @@ -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, diff --git a/tests/session/test_compressor_v3.py b/tests/session/test_compressor_v3.py index 7f8da7ca5c..3ba967e3ed 100644 --- a/tests/session/test_compressor_v3.py +++ b/tests/session/test_compressor_v3.py @@ -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 = []