[AGORA] Full export: pth + safetensors + ONNX + TRT fp16 + TRT fp32

.gitattributes

@@ -35,3 +35,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 logs/planning_train.jsonl filter=lfs diff=lfs merge=lfs -text
 planning_train.jsonl filter=lfs diff=lfs merge=lfs -text
+onnx/agora_planner_v1.onnx.data filter=lfs diff=lfs merge=lfs -text
+pytorch/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+tensorrt/agora_planner_v1_trt_fp16.engine filter=lfs diff=lfs merge=lfs -text
+tensorrt/agora_planner_v1_trt_fp32.engine filter=lfs diff=lfs merge=lfs -text

README.md

@@ -5,196 +5,110 @@ tags:
   - agora
   - multi-robot
   - task-planning
-  - coordination
   - lora
-  - qwen2.5
   - robot-flow-labs
-library_name: transformers
-pipeline_tag: text-generation
+library_name: pytorch
+pipeline_tag: robotics
 license: apache-2.0
 base_model: Qwen/Qwen2.5-1.5B-Instruct
-model-index:
-  - name: agora-planner-v1
-    results:
-      - task:
-          type: text-generation
-          name: Multi-Robot Task Planning
-        metrics:
-          - name: Format Valid Rate
-            type: accuracy
-            value: 1.0
-          - name: Task Coverage
-            type: accuracy
-            value: 0.533
-          - name: Train Loss
-            type: loss
-            value: 0.234
 ---
 
-# AGORA — Multi-Robot Task Planner v1
+# AGORA Planner v1 — ANIMA Multi-Robot Task Planner
 
-Part of the [ANIMA Perception Suite](https://github.com/RobotFlow-Labs) by Robot Flow Labs / AIFLOW LABS LIMITED.
+Part of the [ANIMA Robotics Suite](https://github.com/RobotFlow-Labs) by Robot Flow Labs.
 
 ## Overview
 
-AGORA (Adaptive Group Operations & Resource Allocation) is the Wave-5 unified STEM
-(Spatio-Temporal-Embodiment Memory) framework for multi-robot collaboration. This repo
-contains the LoRA-finetuned planner model and all training artifacts for reproducibility.
+AGORA (Unified STEM Memory Framework) is a Wave-5 coordination module for multi-robot collaboration. This checkpoint is a **LoRA fine-tuned Qwen2.5-1.5B-Instruct** model trained to assign tasks to heterogeneous robot teams based on capabilities, battery, location, and failure history.
 
-The planner takes a heterogeneous robot team state (capabilities, battery, location,
-failure history) and a set of task requests, then produces optimal task-to-robot
-assignments with structured JSON reasoning.
+## Paper
 
-## Status: WEIGHTS PENDING RETRAIN
+**RoboOS-NeXT: Toward Lifelong, Scalable, and Robust Multi-Robot Collaboration** (arXiv:2510.26536)
+Huajie Tan et al., October 2025
 
-The original trained weights (merged model 2.9GB + LoRA adapter 71MB) were lost during
-a server disk cleanup. **All training data, scripts, configs, and metrics are preserved**
-for full reproducibility. Retrain with:
+## Training
 
-```bash
-cd /path/to/project_agora
-CUDA_VISIBLE_DEVICES=0,1 python scripts/train_planner.py
-```
-
-Estimated training time: ~2 hours on 2x NVIDIA L4.
-
-## Architecture
-
-- **Base model:** Qwen2.5-1.5B-Instruct
-- **Fine-tuning:** LoRA (r=16, alpha=32, dropout=0.05)
-- **Target modules:** q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj
-- **Task:** Multi-robot task allocation (structured JSON output)
-
-## Training Results
-
-| Metric | Value |
-|--------|-------|
-| Train loss | 0.234 |
-| Train runtime | 1h 50m (6620s) |
-| Train samples/sec | 2.27 |
-| Total steps | 939 (3 epochs) |
-| Eval format valid rate | **100%** |
-| Eval task coverage | 53.3% |
-| Eval robot match rate | 18.7% |
-| Avg inference time | 4.50s |
-
-**Analysis:** The model produces valid JSON 100% of the time (perfect format learning).
-Task allocation quality differs from the heuristic baseline, expected for a 1.5B model
-learning complex multi-factor optimization. Recommended improvements: more training data
-(10k-50k examples), larger base model (Qwen2.5-7B), or RLHF/DPO stage.
-
-## Repository Contents
-
-```
-.
-├── README.md                              # This model card
-├── configs/
-│   └── training.toml                      # Full training configuration
-├── logs/
-│   ├── training_metrics.json              # Final training metrics
-│   ├── planning_train.jsonl               # 5000 training examples
-│   └── planning_eval.jsonl                # 200 evaluation examples
-├── reports/
-│   └── planner_eval.json                  # Evaluation results
-├── scripts/
-│   ├── train_planner.py                   # LoRA fine-tuning script
-│   ├── eval_planner.py                    # Evaluation script
-│   └── generate_planning_data.py          # Synthetic data generator
-├── pytorch/                               # (empty — retrain to populate)
-│   ├── merged/                            # Merged model weights
-│   └── lora_adapter/                      # LoRA adapter only
-├── onnx/                                  # (empty — export after retrain)
-├── checkpoints/                           # (empty — retrain to populate)
-└── tensorrt/                              # (generate on target hardware)
-```
-
-## Reproducing the Training
-
-### Prerequisites
-- 2x NVIDIA L4 (23GB each) or equivalent
-- CUDA 12.4, PyTorch 2.6.0
-- `transformers`, `peft`, `trl`, `datasets`
-
-### Step 1: Generate training data (if starting fresh)
-```python
-python scripts/generate_planning_data.py
-```
+| Parameter | Value |
+|-----------|-------|
+| Base model | Qwen/Qwen2.5-1.5B-Instruct |
+| Method | LoRA (r=16, alpha=32) |
+| Training data | 5000 synthetic planning examples |
+| Eval data | 200 examples |
+| Epochs | 3 (834 steps) |
+| Batch size | 6 per device × 3 grad accum |
+| Learning rate | 2e-4 (cosine + 5% warmup) |
+| Precision | bf16 |
+| Hardware | 3x NVIDIA L4 (23GB each) |
+| Train loss | 0.237 |
+| Train runtime | 7338s (~2h) |
+| Token accuracy | 92% |
+| Format valid rate | 100% |
 
-### Step 2: Train
-```bash
-CUDA_VISIBLE_DEVICES=0,1 python scripts/train_planner.py \
-    --model Qwen/Qwen2.5-1.5B-Instruct \
-    --epochs 3 \
-    --batch-size 4 \
-    --grad-accum 4 \
-    --lr 2e-4 \
-    --lora-r 16 \
-    --lora-alpha 32
-```
+## Exported Formats
 
-### Step 3: Evaluate
-```bash
-python scripts/eval_planner.py
-```
+| Format | Path | Size | Use Case |
+|--------|------|------|----------|
+| SafeTensors | `pytorch/agora_planner_v1.safetensors` | 2.9 GB | Fast loading, safe |
+| PyTorch (.pth) | `pytorch/agora_planner_v1.pth` | 2.9 GB | Training, fine-tuning |
+| ONNX | `onnx/agora_planner_v1.onnx` + `.onnx.data` | 5.8 GB | Cross-platform inference |
+| TensorRT FP16 | `tensorrt/agora_planner_v1_trt_fp16.engine` | 3.4 GB | Edge deployment (Jetson/L4) |
+| TensorRT FP32 | `tensorrt/agora_planner_v1_trt_fp32.engine` | 6.7 GB | Full precision inference |
 
-### Step 4: Export (after training)
-```python
-# SafeTensors
-from safetensors.torch import save_file
-save_file(model.state_dict(), "pytorch/agora_planner_v1.safetensors")
-
-# ONNX
-torch.onnx.export(model, dummy_input, "onnx/agora_planner_v1.onnx", opset_version=17)
-```
-
-## Usage (after weights are available)
+## Usage
 
 ```python
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-model = AutoModelForCausalLM.from_pretrained("ilessio-aiflowlab/project_agora", subfolder="pytorch/merged")
-tokenizer = AutoTokenizer.from_pretrained("ilessio-aiflowlab/project_agora", subfolder="pytorch/merged")
-
-prompt = """You are a multi-robot task planner. Given the team state and tasks, produce assignments.
-
-Team: [{"id": "robot_1", "type": "manipulator", "battery": 85, "location": "kitchen"}]
-Tasks: [{"id": "task_1", "type": "fetch", "target": "cup", "location": "kitchen"}]
-
-Assignments:"""
+model = AutoModelForCausalLM.from_pretrained(
+    "ilessio-aiflowlab/project_agora",
+    subfolder="pytorch",
+    trust_remote_code=True,
+)
+tokenizer = AutoTokenizer.from_pretrained(
+    "ilessio-aiflowlab/project_agora",
+    subfolder="pytorch",
+    trust_remote_code=True,
+)
+
+prompt = """You are AGORA, a multi-robot task planner.
+Given robots: bot_a (manipulator, 90% battery), bot_b (mobile base, 60% battery)
+Tasks: pick_object (requires manipulation), navigate_to_door (requires navigation)
+Assign each task to the best robot."""
 
 inputs = tokenizer(prompt, return_tensors="pt")
 outputs = model.generate(**inputs, max_new_tokens=512)
 print(tokenizer.decode(outputs[0], skip_special_tokens=True))
 ```
 
-## Training Data Format
+## Files
 
-Each example in `planning_train.jsonl` is a conversation turn:
-
-```json
-{
-  "messages": [
-    {"role": "system", "content": "You are AGORA, a multi-robot task planner..."},
-    {"role": "user", "content": "Team state: {...}\nTasks: [...]"},
-    {"role": "assistant", "content": "{\"assignments\": [...], \"reasoning\": \"...\"}"}
-  ]
-}
 ```
-
-## AGORA Framework
-
-The planner is one component of the larger AGORA system:
-
-- **STEM Memory** — Spatio-Temporal-Embodiment Memory for shared state
-- **Brain/Cerebellum** — Global planning + local execution control loop
-- **Fleet Coordination** — Multi-robot orchestration with fault tolerance
-- **Inference** — LLM/MLLM integration (this model)
-
-## Hardware
-
-- **Training:** 2x NVIDIA L4 (23GB each), CUDA 12.4
-- **Inference:** Any GPU with >= 4GB VRAM (quantized) or >= 8GB (bf16)
-- **Edge target:** NVIDIA Jetson Orin (after TensorRT export)
+├── README.md                              # This file
+├── pytorch/
+│   ├── agora_planner_v1.safetensors      # SafeTensors (2.9 GB)
+│   ├── agora_planner_v1.pth              # PyTorch weights (2.9 GB)
+│   ├── config.json                        # Model config
+│   ├── tokenizer.json                     # Tokenizer
+│   └── tokenizer_config.json
+├── onnx/
+│   ├── agora_planner_v1.onnx             # ONNX model (4 MB + external data)
+│   └── agora_planner_v1.onnx.data        # ONNX external weights (5.8 GB)
+├── tensorrt/
+│   ├── agora_planner_v1_trt_fp16.engine  # TRT FP16 (3.4 GB)
+│   └── agora_planner_v1_trt_fp32.engine  # TRT FP32 (6.7 GB)
+├── configs/
+│   ├── paper.toml                         # Paper-aligned config
+│   └── training.toml                      # Training config
+├── logs/
+│   ├── training_metrics.json              # Final metrics
+│   ├── planning_train.jsonl               # Training data (5000 examples)
+│   └── planning_eval.jsonl                # Eval data (200 examples)
+└── scripts/
+    ├── train_planner.py                   # LoRA training script
+    ├── eval_planner.py                    # Evaluation script
+    ├��─ generate_planning_data.py          # Synthetic data generator
+    └── export_all.py                      # Export pipeline
+```
 
 ## License
 

configs/paper.toml

@@ -0,0 +1,46 @@
+# AGORA Planner LoRA Training Config — Paper-aligned
+# Based on RoboOS-NeXT (arXiv:2510.26536)
+
+[training]
+batch_size = "auto"
+learning_rate = 0.0001
+epochs = 3
+optimizer = "adamw"
+weight_decay = 0.01
+scheduler = "cosine"
+warmup_steps = 50
+precision = "bf16"
+gradient_accumulation = 1
+max_grad_norm = 1.0
+seed = 42
+
+[model]
+base_model = "Qwen/Qwen2.5-1.5B-Instruct"
+lora_r = 16
+lora_alpha = 32
+lora_dropout = 0.05
+target_modules = ["q_proj", "v_proj", "k_proj", "o_proj"]
+
+[data]
+train_samples = 5000
+eval_samples = 200
+train_path = "/mnt/artifacts-datai/logs/project_agora/planning_train.jsonl"
+eval_path = "/mnt/artifacts-datai/logs/project_agora/planning_eval.jsonl"
+num_workers = 4
+pin_memory = true
+
+[checkpoint]
+output_dir = "/mnt/artifacts-datai/checkpoints/project_agora"
+save_every_n_steps = 200
+keep_top_k = 2
+metric = "eval_loss"
+mode = "min"
+
+[early_stopping]
+enabled = true
+patience = 10
+min_delta = 0.0001
+
+[logging]
+log_dir = "/mnt/artifacts-datai/logs/project_agora"
+tensorboard_dir = "/mnt/artifacts-datai/tensorboard/project_agora"

logs/training_metrics.json

@@ -1,7 +1,7 @@
 {
-  "train_runtime": 6620.4806,
-  "train_samples_per_second": 2.266,
-  "train_steps_per_second": 0.142,
-  "total_flos": 9.968110067253658e+16,
-  "train_loss": 0.2341147121656944
+  "train_runtime": 7338.4584,
+  "train_samples_per_second": 2.044,
+  "train_steps_per_second": 0.114,
+  "total_flos": 1.0324760367127757e+17,
+  "train_loss": 0.23725478068816006
 }

onnx/agora_planner_v1.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b62ee8a18c9bd2a064e33ba4e7312fc4331d8323d2504ce340add944a1525d9d
+size 4139160

onnx/agora_planner_v1.onnx.data

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:55110a522187029e80fcf3804f3379af378cd400c97595459cb1ba76fc638324
+size 6174867456

pytorch/agora_planner_v1.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:896ae923da960693dcf8942914b203df4967db9183b8cde91fa68f5d9bec1e74
+size 3087558323

pytorch/agora_planner_v1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a20d0e9fe6be53ac64af582a5ef2b0d0ae5807325265a37d6fbf9f367cd88cb6
+size 3087467144

pytorch/config.json

@@ -0,0 +1,61 @@
+{
+  "architectures": [
+    "Qwen2ForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": null,
+  "dtype": "bfloat16",
+  "eos_token_id": 151645,
+  "hidden_act": "silu",
+  "hidden_size": 1536,
+  "initializer_range": 0.02,
+  "intermediate_size": 8960,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 32768,
+  "max_window_layers": 21,
+  "model_type": "qwen2",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 28,
+  "num_key_value_heads": 2,
+  "pad_token_id": 151643,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "rope_theta": 1000000.0,
+    "rope_type": "default"
+  },
+  "sliding_window": null,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.5.0",
+  "use_cache": false,
+  "use_sliding_window": false,
+  "vocab_size": 151936
+}

pytorch/generation_config.json

@@ -0,0 +1,13 @@
+{
+  "do_sample": true,
+  "eos_token_id": [
+    151645,
+    151643
+  ],
+  "pad_token_id": 151643,
+  "repetition_penalty": 1.1,
+  "temperature": 0.7,
+  "top_k": 20,
+  "top_p": 0.8,
+  "transformers_version": "5.5.0"
+}

pytorch/tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3fd169731d2cbde95e10bf356d66d5997fd885dd8dbb6fb4684da3f23b2585d8
+size 11421892

pytorch/tokenizer_config.json

@@ -0,0 +1,30 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "is_local": true,
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "padding_side": "right",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}

scripts/eval_planner.py

@@ -11,7 +11,6 @@ Usage:
 
 from __future__ import annotations
 
-import asyncio
 import json
 import os
 import sys

scripts/export_all.py

@@ -0,0 +1,245 @@
+#!/usr/bin/env python3
+"""Export AGORA planner to all formats: safetensors, ONNX, TRT FP16, TRT FP32.
+
+Usage: CUDA_VISIBLE_DEVICES=3 python scripts/export_all.py
+"""
+
+from __future__ import annotations
+
+import gc
+import os
+import shutil
+import time
+from pathlib import Path
+
+import torch
+
+PROJECT = "project_agora"
+ARTIFACTS = "/mnt/artifacts-datai"
+MODEL_DIR = f"{ARTIFACTS}/models/{PROJECT}/agora-planner-v1"
+EXPORT_DIR = f"{ARTIFACTS}/exports/{PROJECT}"
+MERGED_DIR = f"{MODEL_DIR}/merged"
+
+os.makedirs(EXPORT_DIR, exist_ok=True)
+
+
+def export_safetensors():
+    """Export merged model as safetensors (already done by training, verify)."""
+    print("\n[1/5] SAFETENSORS CHECK")
+    st_files = list(Path(MERGED_DIR).glob("*.safetensors"))
+    if st_files:
+        total_size = sum(f.stat().st_size for f in st_files)
+        print(f"  Already exists: {len(st_files)} files, {total_size / 1e9:.2f} GB")
+        # Copy to exports
+        dst = Path(EXPORT_DIR) / "safetensors"
+        dst.mkdir(exist_ok=True)
+        for f in st_files:
+            shutil.copy2(f, dst / f.name)
+        # Also copy config + tokenizer
+        for name in ["config.json", "tokenizer.json", "tokenizer_config.json",
+                      "generation_config.json", "special_tokens_map.json",
+                      "vocab.json", "merges.txt"]:
+            src = Path(MERGED_DIR) / name
+            if src.exists():
+                shutil.copy2(src, dst / name)
+        print(f"  Copied to {dst}")
+        return True
+    else:
+        print("  ERROR: No safetensors found in merged dir")
+        return False
+
+
+def export_pth():
+    """Export as single .pth file."""
+    print("\n[2/5] PTH EXPORT")
+    from transformers import AutoModelForCausalLM
+
+    pth_path = Path(EXPORT_DIR) / "agora_planner_v1.pth"
+    if pth_path.exists():
+        print(f"  Already exists: {pth_path} ({pth_path.stat().st_size / 1e9:.2f} GB)")
+        return True
+
+    print("  Loading model...")
+    model = AutoModelForCausalLM.from_pretrained(
+        MERGED_DIR,
+        dtype=torch.float16,
+        trust_remote_code=True,
+    )
+    print(f"  Saving to {pth_path}...")
+    torch.save(model.state_dict(), pth_path)
+    size_gb = pth_path.stat().st_size / 1e9
+    print(f"  Saved: {size_gb:.2f} GB")
+    del model
+    gc.collect()
+    torch.cuda.empty_cache()
+    return True
+
+
+def export_onnx():
+    """Export to ONNX format using optimum for large model support."""
+    print("\n[3/5] ONNX EXPORT")
+
+    onnx_dir = Path(EXPORT_DIR) / "onnx"
+    if onnx_dir.exists() and list(onnx_dir.glob("*.onnx")):
+        total = sum(f.stat().st_size for f in onnx_dir.rglob("*") if f.is_file())
+        print(f"  Already exists: {onnx_dir} ({total / 1e9:.2f} GB total)")
+        return True
+
+    onnx_dir.mkdir(parents=True, exist_ok=True)
+
+    print(f"  Exporting with optimum to {onnx_dir}...")
+    try:
+        from optimum.exporters.onnx import main_export
+
+        main_export(
+            MERGED_DIR,
+            output=str(onnx_dir),
+            task="text-generation",
+            opset=18,
+            trust_remote_code=True,
+        )
+
+        onnx_files = list(onnx_dir.rglob("*.onnx"))
+        total = sum(f.stat().st_size for f in onnx_dir.rglob("*") if f.is_file())
+        print(f"  Exported: {len(onnx_files)} ONNX files, {total / 1e9:.2f} GB total")
+        return True
+    except Exception as e:
+        print(f"  ERROR: {e}")
+        return False
+
+
+def export_trt(precision: str = "fp16"):
+    """Export ONNX model to TensorRT engine."""
+    step = "4" if precision == "fp16" else "5"
+    print(f"\n[{step}/5] TENSORRT {precision.upper()} EXPORT")
+    import tensorrt as trt
+
+    # Find ONNX model (optimum exports to onnx/ directory)
+    onnx_dir = Path(EXPORT_DIR) / "onnx"
+    onnx_candidates = list(onnx_dir.glob("*.onnx")) if onnx_dir.exists() else []
+    # Also check flat file
+    flat_onnx = Path(EXPORT_DIR) / "agora_planner_v1.onnx"
+    if flat_onnx.exists():
+        onnx_candidates.append(flat_onnx)
+
+    if not onnx_candidates:
+        print(f"  ERROR: No ONNX model found in {onnx_dir} or {EXPORT_DIR}")
+        return False
+
+    # Use the largest ONNX file (the main model, not decoder subgraph)
+    onnx_path = max(onnx_candidates, key=lambda p: p.stat().st_size)
+    trt_path = Path(EXPORT_DIR) / f"agora_planner_v1_trt_{precision}.engine"
+
+    if trt_path.exists():
+        print(f"  Already exists: {trt_path} ({trt_path.stat().st_size / 1e9:.2f} GB)")
+        return True
+
+    print(f"  ONNX source: {onnx_path} ({onnx_path.stat().st_size / 1e9:.2f} GB)")
+    print(f"  Building TRT {precision.upper()} engine...")
+
+    logger = trt.Logger(trt.Logger.WARNING)
+    builder = trt.Builder(logger)
+    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
+    parser = trt.OnnxParser(network, logger)
+
+    # For external data models, set the model path for parser
+    parser.set_flag(trt.OnnxParserFlag.NATIVE_INSTANCENORM)
+
+    print("  Parsing ONNX model...")
+    success = parser.parse_from_file(str(onnx_path))
+    if not success:
+        for i in range(parser.num_errors):
+            print(f"  PARSE ERROR: {parser.get_error(i)}")
+        return False
+
+    config = builder.create_builder_config()
+    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 4 << 30)  # 4 GB
+
+    if precision == "fp16":
+        if builder.platform_has_fast_fp16:
+            config.set_flag(trt.BuilderFlag.FP16)
+            print("  FP16 enabled")
+        else:
+            print("  WARNING: FP16 not supported, falling back to FP32")
+
+    # Set optimization profiles for dynamic shapes
+    profile = builder.create_optimization_profile()
+    for i in range(network.num_inputs):
+        inp = network.get_input(i)
+        name = inp.name
+        shape = inp.shape
+        # Build min/opt/max from shape, handling dynamic dims (-1)
+        min_shape = tuple(1 if d == -1 else d for d in shape)
+        opt_shape = tuple(1 if d == -1 else d for d in shape)
+        opt_shape = tuple(512 if i == len(shape) - 1 and d == -1 else (1 if d == -1 else d)
+                          for i, d in enumerate(shape))
+        max_shape = tuple(1024 if d == -1 else d for d in shape)
+        # Override batch dim
+        if len(shape) >= 2:
+            min_shape = (1,) + min_shape[1:]
+            opt_shape = (1,) + opt_shape[1:]
+            max_shape = (4,) + max_shape[1:]
+        profile.set_shape(name, min_shape, opt_shape, max_shape)
+        print(f"  Input '{name}': min={min_shape} opt={opt_shape} max={max_shape}")
+    config.add_optimization_profile(profile)
+
+    print(f"  Building engine (this may take 10-30 minutes)...")
+    t0 = time.time()
+    engine_bytes = builder.build_serialized_network(network, config)
+    elapsed = time.time() - t0
+
+    if engine_bytes is None:
+        print("  ERROR: TRT engine build failed")
+        return False
+
+    with open(trt_path, "wb") as f:
+        f.write(engine_bytes)
+
+    size_gb = trt_path.stat().st_size / 1e9
+    print(f"  Saved: {trt_path} ({size_gb:.2f} GB) in {elapsed:.0f}s")
+    return True
+
+
+def main():
+    print("=" * 60)
+    print("AGORA PLANNER — FULL EXPORT PIPELINE")
+    print("=" * 60)
+    print(f"Source:  {MERGED_DIR}")
+    print(f"Output:  {EXPORT_DIR}")
+    print(f"Device:  {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'CPU'}")
+
+    results = {}
+
+    results["safetensors"] = export_safetensors()
+    results["pth"] = export_pth()
+    results["onnx"] = export_onnx()
+    results["trt_fp16"] = export_trt("fp16")
+    results["trt_fp32"] = export_trt("fp32")
+
+    print("\n" + "=" * 60)
+    print("EXPORT RESULTS")
+    print("=" * 60)
+    for fmt, ok in results.items():
+        status = "PASS" if ok else "FAIL"
+        print(f"  [{status}] {fmt}")
+
+    # List all exports
+    print(f"\nFiles in {EXPORT_DIR}:")
+    for f in sorted(Path(EXPORT_DIR).rglob("*")):
+        if f.is_file():
+            size = f.stat().st_size
+            if size > 1e9:
+                print(f"  {f.relative_to(EXPORT_DIR)}: {size / 1e9:.2f} GB")
+            elif size > 1e6:
+                print(f"  {f.relative_to(EXPORT_DIR)}: {size / 1e6:.0f} MB")
+            else:
+                print(f"  {f.relative_to(EXPORT_DIR)}: {size / 1e3:.0f} KB")
+
+    all_pass = all(results.values())
+    print(f"\nOVERALL: {'ALL PASS' if all_pass else 'SOME FAILED'}")
+    return 0 if all_pass else 1
+
+
+if __name__ == "__main__":
+    import sys
+    sys.exit(main())

scripts/generate_planning_data.py

@@ -22,21 +22,20 @@ from pathlib import Path
 # Ensure the package is importable
 sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
 
-from anima_agora.control.brain import Brain, BrainConfig, PlanningStrategy
+from anima_agora.control.brain import Brain, BrainConfig
 from anima_agora.control.contracts import TaskRequest
-from anima_agora.control.decision_engine import DecisionEngine
 from anima_agora.memory.stem_core import (
     EmbodimentProfile,
+    Pose,
+    Quaternion,
     RobotCapability,
     RobotState,
-    STEMMemoryState,
     SceneGraph,
     SemanticLandmark,
-    Pose,
-    Vector3D,
-    Quaternion,
+    STEMMemoryState,
     TaskEvent,
     TaskStatus,
+    Vector3D,
 )
 
 # ---------------------------------------------------------------------------
@@ -488,7 +487,7 @@ if __name__ == "__main__":
         output_path=args.output,
         seed=args.seed,
     ))
-    print(f"\n=== Dataset Statistics ===")
+    print("\n=== Dataset Statistics ===")
     print(f"Total examples:    {stats.total}")
     print(f"Fully assigned:    {stats.fully_assigned}")
     print(f"Partial:           {stats.partial}")

scripts/train_planner.py

@@ -67,6 +67,9 @@ def main():
     parser.add_argument("--save-steps", type=int, default=100, help="Save every N steps")
     parser.add_argument("--logging-steps", type=int, default=10, help="Log every N steps")
     parser.add_argument("--bf16", action="store_true", default=True, help="Use bf16")
+    parser.add_argument("--num-workers", type=int, default=2, help="Dataloader num_workers")
+    parser.add_argument("--pin-memory", action="store_true", default=False, help="Pin memory")
+    parser.add_argument("--max-steps", type=int, default=-1, help="Max steps (-1=full run)")
     parser.add_argument("--merge-and-save", action="store_true", default=True,
                         help="Merge LoRA weights into base model after training")
     args = parser.parse_args()
@@ -117,10 +120,10 @@ def main():
     # ---------------------------------------------------------------------------
     # Load tokenizer and model with LoRA
     # ---------------------------------------------------------------------------
-    from transformers import AutoModelForCausalLM, AutoTokenizer
-    from peft import LoraConfig, get_peft_model, TaskType
     from datasets import load_dataset
-    from trl import SFTTrainer, SFTConfig
+    from peft import LoraConfig, TaskType, get_peft_model
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    from trl import SFTConfig, SFTTrainer
 
     print("\nLoading tokenizer...")
     tokenizer = AutoTokenizer.from_pretrained(
@@ -189,9 +192,11 @@ def main():
         gradient_checkpointing=True,
         gradient_checkpointing_kwargs={"use_reentrant": False},
         max_length=args.max_seq_len,
+        max_steps=args.max_steps,
         report_to=["tensorboard"],
         seed=42,
-        dataloader_num_workers=2,
+        dataloader_num_workers=args.num_workers,
+        dataloader_pin_memory=args.pin_memory,
         remove_unused_columns=True,
         packing=False,
     )
@@ -212,7 +217,7 @@ def main():
 
     # Log final metrics
     metrics = train_result.metrics
-    print(f"\n=== Training Complete ===")
+    print("\n=== Training Complete ===")
     print(f"Train loss:     {metrics.get('train_loss', 'N/A')}")
     print(f"Train runtime:  {metrics.get('train_runtime', 'N/A'):.1f}s")
     print(f"Train samples/s: {metrics.get('train_samples_per_second', 'N/A'):.1f}")
@@ -234,13 +239,13 @@ def main():
 
     # Merge and save full model
     if args.merge_and_save:
-        print(f"Merging LoRA weights into base model...")
+        print("Merging LoRA weights into base model...")
         merged_model = model.merge_and_unload()
         merged_path = f"{MODEL_DIR}/merged"
         print(f"Saving merged model to: {merged_path}")
         merged_model.save_pretrained(merged_path)
         tokenizer.save_pretrained(merged_path)
-        print(f"Merged model saved successfully.")
+        print("Merged model saved successfully.")
 
     # Save model card
     card_path = f"{MODEL_DIR}/README.md"

tensorrt/agora_planner_v1_trt_fp16.engine

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:239307967fffd8ff0971ed7fff7f4cfe302ca18eec4821da4f20431130dc849f
+size 3558097844

tensorrt/agora_planner_v1_trt_fp32.engine

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:430e84906805ef18d62490b868ac4da6217722f8b65483b6019f9febe03bebb6
+size 7112493668