Spaces:
Runtime error
Runtime error
File size: 2,612 Bytes
02ff91f | 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 | """
Cross-company transfer learning strategy.
Freeze encoder, fine-tune specialist-selection and mode heads only.
50 episodes for same-domain, not 600.
"""
from __future__ import annotations
import os
from pathlib import Path
class TransferLearningStrategy:
"""
Enables rapid adaptation to new company rosters.
Strategy:
- The encoder already understands task-capability semantics
- Only the specialist-selection and mode heads need updating
- Fine-tune for 50 episodes same-domain (vs 600 from scratch)
"""
def __init__(self, base_model_path: str = "checkpoints/spindleflow_final"):
self.base_model_path = Path(base_model_path)
def fine_tune_for_new_roster(
self,
new_catalog_path: str,
new_company_tasks: list[str],
num_episodes: int = 50,
output_path: str = "checkpoints/fine_tuned",
) -> None:
"""
Fine-tune the base policy for a new company's specialist roster.
Implementation:
1. Load base model (encoder weights frozen)
2. Replace specialist registry with new catalog
3. Run fine-tuning for num_episodes
4. Save fine-tuned model
For hackathon: documented as architecture decision.
Full implementation requires loading the SB3 model and
selectively freezing layers.
"""
print(f"[Transfer] Fine-tuning for new roster: {new_catalog_path}")
print(f"[Transfer] Tasks: {len(new_company_tasks)} company-specific tasks")
print(f"[Transfer] Episodes: {num_episodes} (vs 600 from scratch)")
print(f"[Transfer] Strategy: Encoder frozen, selection+mode heads trainable")
print(f"[Transfer] Estimated time: {num_episodes * 2}s (vs 1200s from scratch)")
print(f"[Transfer] NOTE: Full SB3 layer-freezing implementation pending.")
def freeze_encoder_layers(self, model) -> None:
"""
Freeze the encoder layers of the SB3 RecurrentPPO model.
Only specialist-selection and mode heads remain trainable.
"""
frozen_count = 0
for name, param in model.policy.named_parameters():
if "lstm" not in name and "action_net" not in name:
param.requires_grad = False
frozen_count += 1
print(f"[Transfer] Frozen {frozen_count} parameter groups")
trainable = sum(
p.numel() for p in model.policy.parameters() if p.requires_grad
)
total = sum(p.numel() for p in model.policy.parameters())
print(f"[Transfer] Trainable: {trainable:,} / {total:,} parameters")
|