garvitsachdeva/SpindleFlow-RL

1"""2Cross-company transfer learning strategy.3Freeze encoder, fine-tune specialist-selection and mode heads only.450 episodes for same-domain, not 600.5"""6 7from __future__ import annotations8import os9from pathlib import Path10 11 12class TransferLearningStrategy:13    """14    Enables rapid adaptation to new company rosters.15 16    Strategy:17    - The encoder already understands task-capability semantics18    - Only the specialist-selection and mode heads need updating19    - Fine-tune for 50 episodes same-domain (vs 600 from scratch)20    """21 22    def __init__(self, base_model_path: str = "checkpoints/spindleflow_final"):23        self.base_model_path = Path(base_model_path)24 25    def fine_tune_for_new_roster(26        self,27        new_catalog_path: str,28        new_company_tasks: list[str],29        num_episodes: int = 50,30        output_path: str = "checkpoints/fine_tuned",31    ) -> None:32        """33        Fine-tune the base policy for a new company's specialist roster.34 35        Implementation:36        1. Load base model (encoder weights frozen)37        2. Replace specialist registry with new catalog38        3. Run fine-tuning for num_episodes39        4. Save fine-tuned model40 41        For hackathon: documented as architecture decision.42        Full implementation requires loading the SB3 model and43        selectively freezing layers.44        """45        print(f"[Transfer] Fine-tuning for new roster: {new_catalog_path}")46        print(f"[Transfer] Tasks: {len(new_company_tasks)} company-specific tasks")47        print(f"[Transfer] Episodes: {num_episodes} (vs 600 from scratch)")48        print(f"[Transfer] Strategy: Encoder frozen, selection+mode heads trainable")49        print(f"[Transfer] Estimated time: {num_episodes * 2}s (vs 1200s from scratch)")50        print(f"[Transfer] NOTE: Full SB3 layer-freezing implementation pending.")51 52    def freeze_encoder_layers(self, model) -> None:53        """54        Freeze the encoder layers of the SB3 RecurrentPPO model.55        Only specialist-selection and mode heads remain trainable.56        """57        frozen_count = 058        for name, param in model.policy.named_parameters():59            if "lstm" not in name and "action_net" not in name:60                param.requires_grad = False61                frozen_count += 162        print(f"[Transfer] Frozen {frozen_count} parameter groups")63        trainable = sum(64            p.numel() for p in model.policy.parameters() if p.requires_grad65        )66        total = sum(p.numel() for p in model.policy.parameters())67        print(f"[Transfer] Trainable: {trainable:,} / {total:,} parameters")68