"""StapleBridge training epoch and per-epoch validation. The Full Exact-SB formulation: * finite lead-specific feasible plan support H(x) via ``property_free_hard_plan_support`` (chemistry, geometry, edit budget, sequence identity, protected positions and exact committed-plan completion); * ``q_ref(p|x)`` preserved on the legal plan space then conditioned on the hard mask; * exact finite-support teacher ``q*(p|x) ~ q_ref(p|x) exp(-beta E_T)``; * amortized plan controller ``q_theta(p|x)`` masked and renormalised on the same support; * ``L_plan = KL(q* || q_theta)``; * conditional execution policy trained on canonical-completing trajectories with lead-local positive path weights. ``validate_enabled`` computes the per-epoch validation metrics, including ``q_star_vs_q_theta_kl``, the checkpoint-selection metric. Post-hoc evaluation reporting is not part of this training release. """ from __future__ import annotations import json import random import time from collections import Counter from typing import Any import numpy as np import torch from staplebridge.chemistry.state import StapleState from staplebridge.hydrocarbon.curriculum import build_hydrocarbon_demonstration_path from staplebridge.hydrocarbon.exact_sb_cache import ( energy_only_from_config, resolve_exact_sb_target, ) from staplebridge.hydrocarbon.plan_control import ( HydrocarbonPlanControlConfig, best_committed_plan_trajectory, HydrocarbonPlanHead, completes_committed_plan, controlled_plan_log_probabilities, configured_plan_level_objective, describe_plan, empirical_log_probabilities, hierarchical_plan_ranking_enabled, mask_and_renormalize_plan_log_probabilities, plan_entropy, property_free_hard_plan_support, sample_committed_plan_trajectory, sample_distinct_plans, select_plan_and_trajectory, ) from staplebridge.hydrocarbon.plan_reference import enumerate_legal_plans from staplebridge.hydrocarbon.plan_validation import ( aggregate_q_star_diagnostics, exact_sb_validation_enabled, lead_q_star_diagnostics, ) from staplebridge.hydrocarbon.property_energy import ( HydrocarbonPropertyEnergyConfig, HydrocarbonPropertyScorer, required_original_lead_properties, ) from staplebridge.hydrocarbon.tokenizer import tokenize_sequence from staplebridge.models.control_kernel import state_to_features from staplebridge.training.losses import ( assign_lead_local_positive_weights, compute_weighted_path_loss, ) from staplebridge.training.records import ( beam_decode, summarize, terminal_record, ) from staplebridge.training.trajectory import TransitionStep, WeightedTrajectory from staplebridge.utils.profiling import STAGE_TIMER def plan_mode(plan: Any) -> str: return f"{plan.ordered_pair}/i,i+{plan.spacing}" def legal_plans(lead: Any, stack: dict[str, Any]) -> list[Any]: return enumerate_legal_plans( tokenize_sequence(lead.linear_sequence), stack["catalog"], protected_positions=lead.protected_positions, filters=stack["sampler"].filters, ) def property_free_supported_plan_view( lead: Any, initial_state: StapleState, stack: dict[str, Any], config: dict[str, Any], ) -> dict[str, Any]: """Build the one hard-constrained coarse-plan support used everywhere.""" all_plans = legal_plans(lead, stack) context = { "protected_positions": lead.protected_positions, "peptide_ca": (lead.target_context or {}).get("peptide_ca"), } all_reference_weights = stack["sampler"].plan_selection_weights( initial_state, all_plans, context ) support_indices, verdicts = property_free_hard_plan_support( initial_state=initial_state, lead=lead, plans=all_plans, catalog=stack["catalog"], catalog_index=stack["catalog_index"], geometry=stack["geometry"], config=config, ) return { "all_plans": all_plans, "all_reference_weights": all_reference_weights, "support_indices": support_indices, "plans": [all_plans[index] for index in support_indices], # Preserve q_ref on the original legal space, then condition it on the # hard mask. Do not recompute empirical within-mode counts afterward. "reference_weights": [ all_reference_weights[index] for index in support_indices ], "verdicts": verdicts, "context": context, } def plan_key(plan: Any) -> tuple[Any, ...]: return (plan.ordered_pair, plan.spacing, tuple(plan.anchor_pair), plan.block_id) def encode_lead(policy: Any, kernel: Any, state: StapleState) -> torch.Tensor: features = state_to_features(state, kernel.block_to_idx) device = kernel.device features = {key: value.to(device) for key, value in features.items()} with torch.no_grad(): return policy.encoder(**features).detach() def selected_top1_rates(selected: list[dict[str, Any]]) -> dict[str, float]: n = max(len(selected), 1) return { "top1_chemistry_valid_rate": sum(bool(row["chemistry_valid"]) for row in selected) / n, "top1_stapled_rate": sum(bool(row["stapled"]) for row in selected) / n, } def train_enabled_epoch( leads: list[Any], config: dict[str, Any], stack: dict[str, Any], energy_fn: Any, policy: Any, kernel: Any, optimizer: Any, parameters: list[torch.nn.Parameter], head: HydrocarbonPlanHead, plan_cfg: HydrocarbonPlanControlConfig, plan_rng: random.Random, epoch: int, exact_sb_cache: Any = None, ) -> tuple[list[dict[str, Any]], dict[str, Any]]: train_cfg = config["training"] horizon = int(train_cfg["horizon"]) chunk_size = int(train_cfg["chunk_size"]) rows_all: list[dict[str, Any]] = [] losses: list[float] = [] path_losses: list[float] = [] plan_losses: list[float] = [] entropies: list[float] = [] coverage: list[float] = [] disagreements: list[bool] = [] exact_reverse_kls: list[float] = [] exact_forward_objectives: list[float] = [] exact_q_star_entropies: list[float] = [] exact_plans_scored = 0 joint_support_enabled = bool( getattr( getattr(energy_fn, "property_cfg", None), "enable_joint_perm_halflife_support", False, ) ) joint_support_rows: list[dict[str, Any]] = [] hard_mask_totals: Counter[str] = Counter() exact_cache_sources: Counter[str] = Counter() controlled_top_modes: Counter[str] = Counter() weighting_totals: Counter[str] = Counter() # --- profiling (opt-in via config profiling.chunk_timing) ------------ # Wall-clock accounting only: every perf_counter below is read outside the # computation it brackets and never feeds a tensor, and no RNG is drawn, so # the trained model, loss and sampling stream are bit-for-bit identical # whether profiling is on or off. # # CUDA attribution: the per-lead GPU stages (plan/qtheta, terminal/PeptiVerse, # fine rollout's ESM2 forward) each END in a blocking device->host transfer # (`.item()` / `.cpu()` / `.tolist()`) that already synchronises at the stage # boundary, so their wall-clock is accurate without an added sync. The three # chunk-level GPU stages (policy log-prob, backward, optimizer) have no such # trailing transfer, so they get an explicit ``torch.cuda.synchronize()`` -- # once per 32-lead chunk, never inside the per-lead or per-step loops. chunk_timing = bool((config.get("profiling") or {}).get("chunk_timing", False)) _cuda_sync = chunk_timing and torch.cuda.is_available() def _sync() -> None: if _cuda_sync: torch.cuda.synchronize() STAGE_KEYS = ( "plan_qtheta", "qstar_target", "fine_rollout", "neighbor_gen", "terminal_pv", "policy_logprob", "backward", "optimizer", ) stage_seconds: dict[str, float] = {key: 0.0 for key in STAGE_KEYS} _profiling_scorer = getattr(energy_fn, "property_scorer", None) _pv_wrapper = getattr(_profiling_scorer, "predictor", None) _esm2_prev = STAGE_TIMER.snapshot() epoch_started = time.perf_counter() total_leads = len(leads) for chunk_start in range(0, len(leads), chunk_size): chunk_leads = leads[chunk_start : chunk_start + chunk_size] weighted_by_lead: list[list[WeightedTrajectory]] = [] chunk_rows: list[dict[str, Any]] = [] chunk_plan_losses: list[torch.Tensor] = [] chunk_t: dict[str, float] = {key: 0.0 for key in STAGE_KEYS} chunk_counts = { "plans": 0, "legal_plans_before_hard_mask": 0, "hard_masked_plans": 0, "leads_without_hard_support": 0, "trajectories": 0, "transitions": 0, "neighbors": 0, } _cache_hits0 = exact_sb_cache.stats.hits if exact_sb_cache is not None else 0 _cache_misses0 = exact_sb_cache.stats.misses if exact_sb_cache is not None else 0 _pv_misses0 = _pv_wrapper.cache_misses if _pv_wrapper is not None else 0 _pv_hits0 = _pv_wrapper.cache_hits if _pv_wrapper is not None else 0 chunk_started = time.perf_counter() for lead_index, lead in enumerate(chunk_leads, start=chunk_start): lead_weighted: list[WeightedTrajectory] = [] z0 = StapleState(sequence_tokens=tokenize_sequence(lead.linear_sequence)) support_view = property_free_supported_plan_view( lead, z0, stack, config ) all_plans = support_view["all_plans"] support_indices = support_view["support_indices"] plans = support_view["plans"] ref_weights = support_view["reference_weights"] context = support_view["context"] chunk_counts["legal_plans_before_hard_mask"] += len(all_plans) chunk_counts["plans"] += len(plans) chunk_counts["hard_masked_plans"] += len(all_plans) - len(plans) hard_mask_totals["legal_plans_before_hard_mask"] += len(all_plans) hard_mask_totals["hard_supported_plans"] += len(plans) hard_mask_totals["hard_masked_plans"] += len(all_plans) - len(plans) if not plans: chunk_counts["leads_without_hard_support"] += 1 hard_mask_totals["leads_without_hard_support"] += 1 if joint_support_enabled: joint_support_rows.append( { "example_id": str(lead.example_id), "n_hard_supported_plans": 0, "joint_plan_count": 0, "joint_nonempty": False, "joint_fallback": False, "q_star_support_size": 0, "status": "no_hard_support", } ) weighted_by_lead.append(lead_weighted) continue sampled = sample_distinct_plans( plans, ref_weights, plan_cfg.plans_per_lead, plan_rng ) if not sampled: raise RuntimeError(f"no sampled legal plan for {lead.example_id}") key_to_index = {plan_key(plan): index for index, plan in enumerate(plans)} _t = time.perf_counter() lead_embedding = encode_lead(policy, kernel, z0) all_controlled_logp = controlled_plan_log_probabilities( head, lead_embedding, all_plans, support_view["all_reference_weights"], len(z0.sequence_tokens), sequence_tokens=z0.sequence_tokens, peptide_ca=context["peptide_ca"], ) masked_controlled_logp = mask_and_renormalize_plan_log_probabilities( all_controlled_logp, support_indices ) controlled_logp = masked_controlled_logp[ torch.tensor( support_indices, dtype=torch.long, device=masked_controlled_logp.device, ) ] chunk_t["plan_qtheta"] += time.perf_counter() - _t joint_train_row: dict[str, Any] | None = None if plan_cfg.exact_sb_objective: reference_logp = empirical_log_probabilities( ref_weights, controlled_logp.device ) # Deterministic half of the target (plan order, q_ref, terminal # energies, log q*) may come from the persistent cache. q_theta # is always the live `controlled_logp` computed above. _t = time.perf_counter() exact_energy_tensor, _, cache_info = resolve_exact_sb_target( lead=lead, plans=plans, reference_log_probabilities=reference_logp, beta=plan_cfg.exact_sb_beta, energy_fn=energy_fn, initial_state=z0, build_terminal=lambda state, plan: build_hydrocarbon_demonstration_path( state, plan, stack["catalog"] )[-1], cache=exact_sb_cache, energy_only=energy_only_from_config(config), ) exact_cache_sources[str(cache_info["source"])] += 1 target_support_mask = ( torch.tensor( cache_info["target_support_mask"], dtype=torch.bool, device=controlled_logp.device, ) if cache_info.get("target_support_mask") is not None else None ) if joint_support_enabled: joint_train_row = { "example_id": str(lead.example_id), "n_hard_supported_plans": len(plans), "joint_plan_count": int(cache_info["joint_plan_count"]), "joint_nonempty": bool(cache_info["joint_nonempty"]), "joint_fallback": bool(cache_info["joint_fallback"]), "q_star_support_size": int(cache_info["q_star_support_size"]), "selection_semantics": "q_theta_top1_hard_plan", "status": "scored", } joint_support_rows.append(joint_train_row) chunk_t["qstar_target"] += time.perf_counter() - _t _t = time.perf_counter() exact_loss, forward_objective, log_q_star = configured_plan_level_objective( controlled_logp, [], [], plan_cfg.target_temperature, exact_sb_objective=True, reference_log_probabilities=reference_logp, terminal_energies=exact_energy_tensor, exact_sb_beta=plan_cfg.exact_sb_beta, target_support_mask=target_support_mask, ) chunk_plan_losses.append(exact_loss) exact_reverse_kls.append(float(exact_loss.detach().cpu().item())) exact_forward_objectives.append( float(forward_objective.detach().cpu().item()) ) q_star = log_q_star.exp() if target_support_mask is None: # Preserve historical false/fallback arithmetic exactly. q_star_entropy = -(q_star * log_q_star).sum() else: q_star_entropy = -torch.where( q_star > 0.0, q_star * log_q_star, torch.zeros_like(q_star), ).sum() exact_q_star_entropies.append( float(q_star_entropy.detach().cpu().item()) ) exact_plans_scored += len(plans) chunk_t["plan_qtheta"] += time.perf_counter() - _t reference_top = int(np.argmax(ref_weights)) controlled_top = int(torch.argmax(controlled_logp).item()) if joint_train_row is not None: joint_train_row.update( { "selected_plan_index": controlled_top, "selected_plan": describe_plan(plans[controlled_top]), "selected_satisfies_joint_condition": bool( target_support_mask is not None and target_support_mask[controlled_top].item() ), } ) disagreements.append(reference_top != controlled_top) entropies.append(float(plan_entropy(controlled_logp).detach().cpu().item())) coverage.append(len(sampled) / len(plans)) controlled_top_modes[ f"{plans[controlled_top].ordered_pair}/i,i+{plans[controlled_top].spacing}" ] += 1 sampled_indices: list[int] = [] penetrance_values: list[float] = [] for trajectory_index, plan in enumerate(sampled): _t = time.perf_counter() trajectory = sample_committed_plan_trajectory( stack["sampler"], z0, plan, lead.protected_positions, context, horizon, ) chunk_t["fine_rollout"] += time.perf_counter() - _t chunk_counts["trajectories"] += 1 trajectory_context = {**context, "hydrocarbon_plan": plan} steps: list[TransitionStep] = [] for t, (state, next_state) in enumerate( zip(trajectory.states[:-1], trajectory.states[1:]) ): _t = time.perf_counter() candidates = stack["graph"].neighbors( state, protected_positions=lead.protected_positions ) chunk_t["neighbor_gen"] += time.perf_counter() - _t chunk_counts["neighbors"] += len(candidates) chunk_counts["transitions"] += 1 chosen = next( (index for index, candidate in enumerate(candidates) if candidate == next_state), None, ) if chosen is None: raise RuntimeError("committed transition absent from graph") steps.append(TransitionStep(state, next_state, candidates, chosen, t)) terminal = trajectory.states[-1] _t = time.perf_counter() row = terminal_record( z0, terminal, lead, energy_fn, stack, epoch=epoch, lead_index=lead_index, trajectory_index=trajectory_index, committed_plan=describe_plan(plan), committed_plan_completed=trajectory.progress.plan_completed, path_length=len(steps), ) chunk_t["terminal_pv"] += time.perf_counter() - _t penetrance = row.get( "hydrocarbon_permeability_penetrance_product_mean" ) penetrance_values.append(float(penetrance) if penetrance is not None else -1.0) sampled_indices.append(key_to_index[plan_key(plan)]) exact_completion = completes_committed_plan(terminal, plan) failure_reason = ( None if exact_completion else "off_plan_completion" if terminal.topology == "stapled" else "unfinished" ) row["trajectory_training_status"] = ( "positive" if exact_completion else failure_reason ) lead_weighted.append( WeightedTrajectory( steps=steps, terminal_state=terminal, terminal_energy=float(row["terminal_energy"]), context=trajectory_context, is_positive=exact_completion, failure_reason=failure_reason, ) ) chunk_rows.append(row) weighted_by_lead.append(lead_weighted) if not plan_cfg.exact_sb_objective: _t = time.perf_counter() legacy_loss, _, _ = configured_plan_level_objective( controlled_logp, sampled_indices, penetrance_values, plan_cfg.target_temperature, exact_sb_objective=False, ) chunk_plan_losses.append(legacy_loss) chunk_t["plan_qtheta"] += time.perf_counter() - _t weight_diagnostics = assign_lead_local_positive_weights(weighted_by_lead) for key in ("n_positive", "n_failure", "n_leads_without_positive"): weighting_totals[key] += int(weight_diagnostics[key]) weighted = [ trajectory for group in weighted_by_lead for trajectory in group ] if not chunk_plan_losses: rows_all.extend(chunk_rows) continue _sync() _t = time.perf_counter() log_probs = [ kernel.log_prob_of( step.state, step.candidates, step.chosen_idx, step.t, trajectory.context, ) for trajectory in weighted for step in trajectory.steps ] path_loss = compute_weighted_path_loss(weighted, log_probs) plan_loss = torch.stack(chunk_plan_losses).mean() total_loss = path_loss + float(plan_cfg.loss_weight) * plan_loss _sync() chunk_t["policy_logprob"] += time.perf_counter() - _t optimizer.zero_grad() _t = time.perf_counter() total_loss.backward() _sync() chunk_t["backward"] += time.perf_counter() - _t _t = time.perf_counter() torch.nn.utils.clip_grad_norm_( parameters, float(train_cfg.get("grad_clip_norm", 1.0)) ) optimizer.step() _sync() chunk_t["optimizer"] += time.perf_counter() - _t losses.append(float(total_loss.detach().cpu().item())) path_losses.append(float(path_loss.detach().cpu().item())) plan_losses.append(float(plan_loss.detach().cpu().item())) rows_all.extend(chunk_rows) for key in STAGE_KEYS: stage_seconds[key] += chunk_t[key] if chunk_timing: _esm2_now = STAGE_TIMER.snapshot() esm2_diff = { key: _esm2_now.get(key, 0) - _esm2_prev.get(key, 0) for key in ( "esm2_forward_time", "esm2_cache_hit", "esm2_cache_miss", "esm2_forward_batches", "esm2_forward_calls", ) } _esm2_prev = _esm2_now leads_done = chunk_start + len(chunk_leads) elapsed = time.perf_counter() - epoch_started leads_per_min = leads_done / (elapsed / 60.0) if elapsed > 0 else 0.0 eta_seconds = ( (elapsed / leads_done) * (total_leads - leads_done) if leads_done else 0.0 ) record = { "epoch": epoch + 1, "chunk": chunk_start // chunk_size + 1, "leads_done": leads_done, "total_leads": total_leads, "chunk_seconds": round(time.perf_counter() - chunk_started, 3), "stage_seconds": {key: round(chunk_t[key], 4) for key in STAGE_KEYS}, "esm2": { key: (round(value, 4) if "time" in key else int(value)) for key, value in esm2_diff.items() }, "counts": dict(chunk_counts), "exact_sb_cache": { "hits": (exact_sb_cache.stats.hits - _cache_hits0) if exact_sb_cache is not None else 0, "misses": (exact_sb_cache.stats.misses - _cache_misses0) if exact_sb_cache is not None else 0, }, "peptiverse": { "predictions": (_pv_wrapper.cache_misses - _pv_misses0) if _pv_wrapper is not None else 0, "cache_hits": (_pv_wrapper.cache_hits - _pv_hits0) if _pv_wrapper is not None else 0, }, "leads_per_min": round(leads_per_min, 1), "epoch_eta_min": round(eta_seconds / 60.0, 1), } print(f"[chunk timing] {json.dumps(record, ensure_ascii=False)}", flush=True) metrics = { "loss": float(np.mean(losses)), "path_loss": float(np.mean(path_losses)), "plan_loss": float(np.mean(plan_losses)), "plan_entropy": float(np.mean(entropies)), "plan_coverage": float(np.mean(coverage)), "reference_vs_controlled_plan_disagreement": float(np.mean(disagreements)), "controlled_plan_mode_mix": dict(controlled_top_modes), "exact_sb_objective": bool(plan_cfg.exact_sb_objective), "exact_sb_beta": float(plan_cfg.exact_sb_beta), "plan_encoder_v2": bool(plan_cfg.plan_encoder_v2), "legacy_target_temperature": float(plan_cfg.target_temperature), "legacy_property_only_equivalent_beta": float( 1.0 / ( plan_cfg.target_temperature * float( config["hydrocarbon"]["terminal_energy"]["property"][ "penetrance_weight" ] ) ) ), "q_star_vs_q_theta_kl": ( float(np.mean(exact_reverse_kls)) if exact_reverse_kls else None ), "exact_sb_forward_objective": ( float(np.mean(exact_forward_objectives)) if exact_forward_objectives else None ), "q_star_entropy": ( float(np.mean(exact_q_star_entropies)) if exact_q_star_entropies else None ), "exact_sb_plans_scored": int(exact_plans_scored), "exact_sb_target_sources": dict(exact_cache_sources), "property_free_hard_plan_mask": True, "hard_plan_mask_uses_peptiverse": False, "hard_plan_mask_totals": dict(hard_mask_totals), "trajectory_weighting": { "normalization": "lead_local_positive_only", **{key: int(value) for key, value in weighting_totals.items()}, }, "stage_seconds": {key: float(value) for key, value in stage_seconds.items()}, } if joint_support_enabled: supported_joint_rows = [ row for row in joint_support_rows if row["status"] == "scored" ] solution_leads = sum( bool(row["joint_nonempty"]) for row in supported_joint_rows ) selected_joint = sum( bool(row["selected_satisfies_joint_condition"]) for row in supported_joint_rows ) metrics["joint_perm_halflife_support"] = { "enabled": True, "leads_total": len(joint_support_rows), "leads_with_hard_support": len(supported_joint_rows), "leads_without_hard_support": sum( row["status"] == "no_hard_support" for row in joint_support_rows ), "leads_with_joint_solution": solution_leads, "fallback_leads": sum( bool(row["joint_fallback"]) for row in supported_joint_rows ), "joint_plan_count": sum( int(row["joint_plan_count"]) for row in supported_joint_rows ), "q_star_support_size": sum( int(row["q_star_support_size"]) for row in supported_joint_rows ), "selected_joint_solutions": selected_joint, "joint_solution_recovery": ( selected_joint / solution_leads if solution_leads else None ), "per_lead": joint_support_rows, } return rows_all, metrics @torch.no_grad() def validate_enabled( leads: list[Any], config: dict[str, Any], stack: dict[str, Any], energy_fn: Any, policy: Any, kernel: Any, head: HydrocarbonPlanHead, exact_sb_cache: Any = None, ) -> tuple[list[dict[str, Any]], dict[str, Any]]: validation_cfg = config["validation"] selected: list[dict[str, Any]] = [] entropies: list[float] = [] coverages: list[float] = [] disagreements: list[bool] = [] top1_plans: dict[str, str] = {} mode_mix: Counter[str] = Counter() # Diagnostics-only: opt-in exact-SB q* on the same property-free hard # support used by q_theta. It never changes fine decoding or candidate # ranking within the committed plan selected below. exact_sb_diagnostics = exact_sb_validation_enabled(config) exact_sb_beta = HydrocarbonPlanControlConfig.from_config(config).exact_sb_beta q_star_rows: list[dict[str, Any]] = [] q_star_cache_sources: Counter[str] = Counter() joint_support_enabled = bool( getattr( getattr(energy_fn, "property_cfg", None), "enable_joint_perm_halflife_support", False, ) ) joint_support_rows: list[dict[str, Any]] = [] legal_plans_before_hard_mask = 0 hard_supported_plans = 0 leads_without_hard_support = 0 # Wall-clock accounting only; see the note in train_enabled_epoch. q_star_diagnostics_seconds = 0.0 for lead_index, lead in enumerate(leads): z0 = StapleState(sequence_tokens=tokenize_sequence(lead.linear_sequence)) support_view = property_free_supported_plan_view(lead, z0, stack, config) all_plans = support_view["all_plans"] support_indices = support_view["support_indices"] plans = support_view["plans"] ref_weights = support_view["reference_weights"] context = support_view["context"] legal_plans_before_hard_mask += len(all_plans) hard_supported_plans += len(plans) strict_support_at_lead = bool(plans) if not plans: leads_without_hard_support += 1 if joint_support_enabled: joint_support_rows.append( { "example_id": str(lead.example_id), "n_hard_supported_plans": 0, "joint_plan_count": 0, "joint_nonempty": False, "joint_fallback": False, "q_star_support_size": 0, "selected_satisfies_joint_condition": False, "status": "no_hard_support", } ) print( f"[plan validation] leads={lead_index + 1}/{len(leads)} " "status=no_strict_feasible_plan", flush=True, ) continue lead_embedding = encode_lead(policy, kernel, z0) all_controlled_logp = controlled_plan_log_probabilities( head, lead_embedding, all_plans, support_view["all_reference_weights"], len(z0.sequence_tokens), sequence_tokens=z0.sequence_tokens, peptide_ca=context["peptide_ca"], ) masked_controlled_logp = mask_and_renormalize_plan_log_probabilities( all_controlled_logp, support_indices ) controlled_logp = masked_controlled_logp[ torch.tensor( support_indices, dtype=torch.long, device=masked_controlled_logp.device, ) ] controlled_probabilities = controlled_logp.detach().exp().cpu() plan_modes = [plan_mode(plan) for plan in plans] q_theta_metrics = { "q_theta_entropy": float(plan_entropy(controlled_logp).cpu().item()), "q_theta_top1_mode": plan_modes[int(torch.argmax(controlled_logp).item())], "q_theta_s5_s5_i4_probability_mass": float( sum( controlled_probabilities[index].item() for index, mode in enumerate(plan_modes) if mode == "S5-S5/i,i+4" ) ), "q_theta_r8_s5_i7_probability_mass": float( sum( controlled_probabilities[index].item() for index, mode in enumerate(plan_modes) if mode == "R8-S5/i,i+7" ) ), } entropies.append(float(plan_entropy(controlled_logp).cpu().item())) reference_top = int(np.argmax(ref_weights)) controlled_top = int(torch.argmax(controlled_logp).item()) disagreements.append(reference_top != controlled_top) coverages.append(1.0) lead_q_star_row: dict[str, Any] | None = None if exact_sb_diagnostics: _diagnostics_started = time.perf_counter() reference_logp = empirical_log_probabilities( ref_weights, controlled_logp.device ) # Same deterministic half as training, same cache. q_theta below is # always recomputed from the live head. plan_energy_tensor, _, cache_info = resolve_exact_sb_target( lead=lead, plans=plans, reference_log_probabilities=reference_logp, beta=exact_sb_beta, energy_fn=energy_fn, initial_state=z0, build_terminal=lambda state, plan: build_hydrocarbon_demonstration_path( state, plan, stack["catalog"] )[-1], cache=exact_sb_cache, energy_only=energy_only_from_config(config), ) q_star_cache_sources[str(cache_info["source"])] += 1 target_support_mask = ( torch.tensor( cache_info["target_support_mask"], dtype=torch.bool, device=controlled_logp.device, ) if cache_info.get("target_support_mask") is not None else None ) diagnostics = lead_q_star_diagnostics( controlled_logp, reference_logp, plan_energy_tensor, exact_sb_beta, plan_modes=plan_modes, target_support_mask=target_support_mask, ) lead_q_star_row = { "example_id": str(lead.example_id), "q_star_top1_plan": describe_plan( plans[diagnostics["q_star_top1_index"]] ), "q_theta_top1_plan": describe_plan( plans[diagnostics["q_theta_top1_index"]] ), "joint_perm_halflife_support_enabled": bool( cache_info["joint_perm_halflife_support_enabled"] ), "joint_plan_count": cache_info["joint_plan_count"], "joint_nonempty": cache_info["joint_nonempty"], "joint_fallback": cache_info["joint_fallback"], "q_star_support_size": int(cache_info["q_star_support_size"]), **diagnostics, } q_star_rows.append(lead_q_star_row) q_star_diagnostics_seconds += time.perf_counter() - _diagnostics_started hierarchical = hierarchical_plan_ranking_enabled(config) decoded_by_plan: list[tuple[StapleState, float, int] | None] = [None] * len(plans) decode_indices = [controlled_top] if hierarchical else range(len(plans)) for plan_index in decode_indices: plan = plans[plan_index] decoded = beam_decode( lead, stack, kernel, int(validation_cfg["horizon"]), int(validation_cfg["beam_size"]), committed_plan=plan, ) if decoded: decoded_by_plan[plan_index] = ( best_committed_plan_trajectory(decoded, plan) if hierarchical else decoded[0] ) try: joint, plan_index, terminal, path_logp, path_length = select_plan_and_trajectory( controlled_logp.detach().cpu(), decoded_by_plan, hierarchical=hierarchical ) except RuntimeError as exc: raise RuntimeError(f"no decoded plans for {lead.example_id}") from exc selected_plan = plans[plan_index] label = describe_plan(selected_plan) top1_plans[lead.example_id] = label mode_mix[ f"{selected_plan.ordered_pair}/i,i+{selected_plan.spacing}" ] += 1 selected_record = terminal_record( z0, terminal, lead, energy_fn, stack, lead_index=lead_index, committed_plan=label, plan_log_probability=float(controlled_logp[plan_index].cpu().item()), path_log_probability=path_logp, joint_log_probability=joint, path_length=path_length, ) selected.append(selected_record) print(f"[plan validation] leads={lead_index + 1}/{len(leads)}", flush=True) summary = summarize(selected, selected, len(leads)) summary.update(selected_top1_rates(selected)) summary.update( { "plan_control_enabled": True, "hierarchical_plan_ranking": hierarchical_plan_ranking_enabled(config), "plan_entropy": float(np.mean(entropies)) if entropies else None, "plan_top1": top1_plans, "plan_coverage": float(np.mean(coverages)) if coverages else 0.0, "mode_mix": dict(mode_mix), "reference_vs_controlled_plan_disagreement": ( float(np.mean(disagreements)) if disagreements else None ), "property_free_hard_plan_mask": True, "hard_plan_mask_uses_peptiverse": False, "legal_plans_before_hard_mask": int(legal_plans_before_hard_mask), "hard_supported_plans": int(hard_supported_plans), "hard_masked_plans": int( legal_plans_before_hard_mask - hard_supported_plans ), "leads_without_hard_support": int(leads_without_hard_support), } ) summary["exact_sb_validation_diagnostics"] = exact_sb_diagnostics if exact_sb_diagnostics: # Additive only: no existing validation metric above is overwritten. summary.update(aggregate_q_star_diagnostics(q_star_rows)) summary["q_star_per_lead"] = q_star_rows # Provenance only; not a metric and not used by any selection. summary["q_star_target_sources"] = dict(q_star_cache_sources) summary["q_star_diagnostics_seconds"] = float(q_star_diagnostics_seconds) if joint_support_enabled: supported_joint_rows = [ row for row in joint_support_rows if row["status"] == "scored" ] solution_leads = sum( bool(row["joint_nonempty"]) for row in supported_joint_rows ) selected_joint = sum( bool(row["selected_satisfies_joint_condition"]) for row in supported_joint_rows ) summary["joint_perm_halflife_support"] = { "enabled": True, "leads_total": len(joint_support_rows), "leads_with_hard_support": len(supported_joint_rows), "leads_without_hard_support": sum( row["status"] == "no_hard_support" for row in joint_support_rows ), "leads_with_joint_solution": solution_leads, "fallback_leads": sum( bool(row["joint_fallback"]) for row in supported_joint_rows ), "joint_plan_count": sum( int(row["joint_plan_count"]) for row in supported_joint_rows ), "q_star_support_size": sum( int(row["q_star_support_size"]) for row in supported_joint_rows ), "selected_joint_solutions": selected_joint, "joint_solution_recovery": ( selected_joint / solution_leads if solution_leads else None ), "per_lead": joint_support_rows, } return selected, summary