| from __future__ import annotations |
|
|
| import torch |
|
|
| from staplebridge.training.trajectory import WeightedTrajectory |
|
|
|
|
| def assign_lead_local_positive_weights( |
| trajectory_groups: list[list[WeightedTrajectory]], |
| ) -> dict[str, float | int]: |
| """Assign lead-equal, within-lead Boltzmann weights to positive paths. |
| |
| Exact committed-plan completions are normal behavior-cloning positives. |
| Off-plan and unfinished paths remain explicit failures with zero positive |
| weight; teaching their chosen actions as positives would contradict strict |
| committed-plan inference. A separate failure-avoidance objective may be |
| added later without changing this positive-path definition. |
| """ |
| if not trajectory_groups: |
| return { |
| "n_leads": 0, |
| "n_positive": 0, |
| "n_failure": 0, |
| "n_leads_without_positive": 0, |
| "total_positive_weight": 0.0, |
| } |
| lead_scale = 1.0 / float(len(trajectory_groups)) |
| n_positive = 0 |
| n_failure = 0 |
| n_without_positive = 0 |
| for group in trajectory_groups: |
| for trajectory in group: |
| trajectory.weight = 0.0 |
| positives = [trajectory for trajectory in group if trajectory.is_positive] |
| n_positive += len(positives) |
| n_failure += len(group) - len(positives) |
| if not positives: |
| n_without_positive += 1 |
| continue |
| energies = torch.tensor( |
| [trajectory.terminal_energy for trajectory in positives], |
| dtype=torch.float64, |
| ) |
| weights = torch.softmax(-energies, dim=0).tolist() |
| for trajectory, weight in zip(positives, weights): |
| trajectory.weight = lead_scale * float(weight) |
| return { |
| "n_leads": len(trajectory_groups), |
| "n_positive": n_positive, |
| "n_failure": n_failure, |
| "n_leads_without_positive": n_without_positive, |
| "total_positive_weight": float( |
| sum(trajectory.weight for group in trajectory_groups for trajectory in group) |
| ), |
| } |
|
|
|
|
| def compute_weighted_path_loss(trajs: list[WeightedTrajectory], log_probs: list[torch.Tensor]) -> torch.Tensor: |
| device = log_probs[0].device if log_probs else torch.device("cpu") |
| weights = torch.tensor([t.weight for t in trajs], dtype=torch.float32, device=device) |
| if weights.numel() == 0: |
| return torch.tensor(0.0, requires_grad=True, device=device) |
| total = torch.tensor(0.0, dtype=torch.float32, device=device) |
| offset = 0 |
| for i, traj in enumerate(trajs): |
| n = len(traj.steps) |
| if n == 0: |
| continue |
| lp = torch.stack(log_probs[offset : offset + n]).sum() |
| total = total - weights[i] * lp |
| offset += n |
| return total |
|
|
