File size: 2,743 Bytes
bb6d2aa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | 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
|