vla / FULL_PIPELINE_DETAILED.md
anhtld's picture
Initial commit: DoVLA-CIL codebase (h=16 breakthrough)
adc02fa verified
|
Raw
History Blame
13.2 kB

πŸš€ FULL PIPELINE: 3-Week Detailed Implementation Plan

Goal: 42-44% β†’ 60-70%+ (SOTA-competitive)

Status: Approved for full implementation with unlimited LLM API


πŸ“… WEEK 1: Language & Data (Day 1-7)

Day 1: Language Embeddings Setup

Morning (4h):

# Install dependencies
pip install sentence-transformers openai anthropic

# Test embedding generation
python -c "
from sentence_transformers import SentenceTransformer
model = SentenceTransformer('all-mpnet-base-v2')
emb = model.encode(['pick the cube'])
print(f'Embedding shape: {emb.shape}')  # Should be (1, 768)
"

Afternoon (4h):

  • Create instruction embedding script
  • Generate embeddings for all 3.5K groups
  • Save to disk (cache for fast loading)

Files to create:

  • dovla_cil/utils/language_embeddings.py
  • scripts/generate_instruction_embeddings.py

Day 2: Modify Architecture for Language

Morning (4h):

  • Update DoVLATransformer to accept lang_dim=768
  • Modify cross-attention to fuse obs+lang
  • Test forward/backward with language

Afternoon (4h):

  • Update training dataset to load embeddings
  • Modify collate_fn for language batching
  • Test full training loop

Files to modify:

  • dovla_cil/models/dovla_transformer.py
  • scripts/train_dovla_transformer.py

Day 3-4: Retrain with Language (48h)

Submit 3 jobs:

sbatch scripts/slurm/train_transformer_lang.sbatch  # 3 seeds

Monitor training:

  • Val top-1 should be 65-70% (vs 63% without lang)
  • Losses should decrease smoothly
  • Expected final: 50-55% selected success

While training runs:

  • Prepare LLM data augmentation code
  • Setup OpenClaude API integration

Day 5: LLM Data Augmentation

Morning (4h):

  • OpenClaude API integration
  • Synthetic instruction generation
def generate_synthetic_instructions(state_desc, num=5):
    prompt = f"""
    Given robot state: {state_desc}
    Generate {num} diverse instructions that could be goals.
    Format: one per line, natural language.
    
    Examples:
    - Pick up the red cube
    - Move the cube to the left
    - Stack the blue block on top
    """
    
    response = openai.ChatCompletion.create(
        model="gpt-4",  # Or claude-3-opus
        messages=[{"role": "user", "content": prompt}]
    )
    return response.choices[0].message.content.split('\n')

Afternoon (4h):

  • Generate synthetic data for 10K additional samples
  • Create augmented dataset
  • Validate quality manually (sample 100)

Day 6: Counterfactual Explanations

LLM-generated failure explanations:

def explain_failure(state, action, outcome):
    prompt = f"""
    State: {state}
    Action: {action}
    Outcome: {outcome['success']} (reward: {outcome['reward']})
    
    In 1 sentence, explain why this action 
    {'succeeded' if outcome['success'] else 'failed'}.
    
    Focus on physical reasoning and constraints.
    """
    
    explanation = claude_api.generate(prompt)
    return explanation

Generate explanations for:

  • All 56K actions in dataset
  • Focus on failures (more informative)
  • Cache to disk

Day 7: Retrain with Augmented Data

Submit training with:

  • Original 3.5K groups
  • +10K synthetic instruction variations
  • +56K action explanations (auxiliary loss)

Expected improvement: +2-5% β†’ 52-57% total


πŸ“… WEEK 2: Architecture & Training (Day 8-14)

Day 8-9: Multi-Scale Transformer

Architecture:

class MultiScaleTransformer(nn.Module):
    def __init__(self):
        self.small = DoVLATransformer(d_model=128, n_layers=2)
        self.medium = DoVLATransformer(d_model=256, n_layers=3)
        self.large = DoVLATransformer(d_model=512, n_layers=4)
        
        # Learned ensemble weights
        self.ensemble_weights = nn.Parameter(torch.ones(3))
    
    def forward(self, obs, actions, lang):
        s1 = self.small(obs, actions, lang)
        s2 = self.medium(obs, actions, lang)
        s3 = self.large(obs, actions, lang)
        
        weights = F.softmax(self.ensemble_weights, dim=0)
        return weights[0]*s1 + weights[1]*s2 + weights[2]*s3

Train 3 scales separately, then ensemble


Day 10: Action-Conditioned Attention

Add action-specific attention:

class ActionConditionedAttention(nn.Module):
    def __init__(self):
        # Learn to attend to relevant state parts per action
        self.action_encoder = nn.Linear(action_dim, d_model)
        self.state_attention = nn.MultiheadAttention(d_model, n_heads)
    
    def forward(self, state, action):
        # Action vector guides what to look at in state
        action_query = self.action_encoder(action)
        attended_state, _ = self.state_attention(
            action_query, state, state
        )
        return attended_state

Day 11-12: Hard Negative Mining

Mine confusing pairs:

def mine_hard_negatives(model, dataset, k=5):
    """Find pairs where model is most confused."""
    hard_pairs = []
    
    for group in dataset:
        scores = model.predict(group)
        
        # Find pairs where:
        # 1. Model predicts A > B
        # 2. Ground truth is B > A
        # 3. Margin is small (confusing)
        
        for i, j in all_pairs:
            pred_margin = scores[i] - scores[j]
            true_margin = rewards[i] - rewards[j]
            
            if sign(pred_margin) != sign(true_margin):
                confusion = abs(pred_margin)
                if confusion < threshold:  # Close call
                    hard_pairs.append((group, i, j, confusion))
    
    # Return top-k% hardest
    return sorted(hard_pairs, key=lambda x: x[-1])[:int(len(hard_pairs)*k/100)]

Retrain focusing 70% on hard pairs, 30% on all pairs


Day 13: Curriculum Learning

Task difficulty ranking:

task_difficulty = {
    'PickCube-v1': 1,      # Easy
    'PushCube-v1': 2,      # Medium
    'PullCube-v1': 2,      # Medium
    'LiftPegUpright-v1': 3,  # Hard
    'StackCube-v1': 4,     # Very hard
    'PegInsertionSide-v1': 5  # Hardest
}

# Training schedule
def get_tasks_for_epoch(epoch, total_epochs=50):
    progress = epoch / total_epochs
    max_difficulty = 1 + progress * 4  # 1 β†’ 5 over training
    
    return [t for t, d in task_difficulty.items() if d <= max_difficulty]

Day 14: Self-Training with LLM Feedback

LLM provides corrective feedback:

def get_llm_feedback(state, action_a, action_b, model_pred, ground_truth):
    if model_pred == ground_truth:
        return None  # Model correct
    
    prompt = f"""
    The model incorrectly predicted action A is better than B.
    Actually, B is better.
    
    State: {state}
    Action A: {action_a}
    Action B: {action_b}
    
    What physical reasoning explains why B > A?
    What should the model learn to avoid this mistake?
    
    Response format:
    - Key insight: [1 sentence]
    - Focus on: [state feature to attend to]
    """
    
    feedback = claude_api.generate(prompt)
    return feedback

Use feedback as auxiliary training signal

Week 2 expected result: 57-62%


πŸ“… WEEK 3: Ensemble & Advanced (Day 15-21)

Day 15-16: Multi-Model Ensemble

Train 5 diverse architectures:

models = {
    'transformer_small': DoVLATransformer(d_model=256, n_layers=2),
    'transformer_large': DoVLATransformer(d_model=512, n_layers=4),
    'mlp_deep': DeepMLP(hidden=[512, 512, 256]),
    'multiscale': MultiScaleTransformer(),
    'action_conditioned': ActionConditionedTransformer()
}

# Train each independently
for name, model in models.items():
    train(model, dataset)
    save(model, f'checkpoints/{name}_best.pt')

Ensemble strategies:

  • Voting (majority vote)
  • Averaging (mean scores)
  • Stacking (meta-learner on top)

Day 17-18: LLM as Final Judge

Most powerful improvement (+10-15%):

def llm_action_ranking(state, instruction, candidate_actions, model_scores):
    """Use LLM to re-rank top-k actions from model."""
    
    # Get top-5 from model ensemble
    top_k = 5
    top_actions = get_top_k(candidate_actions, model_scores, k=top_k)
    
    # Format for LLM
    action_descriptions = [
        f"{i+1}. {describe_action(a)}" 
        for i, a in enumerate(top_actions)
    ]
    
    prompt = f"""
    You are a robot action selection expert.
    
    State:
    {describe_state(state)}
    
    Goal:
    {instruction}
    
    Candidate actions:
    {chr(10).join(action_descriptions)}
    
    Rank these actions from 1 (best) to {top_k} (worst).
    Consider:
    - Physics (will it work?)
    - Safety (any collisions?)
    - Efficiency (direct path?)
    - Goal achievement
    
    Output ONLY the ranking numbers: [best_idx, 2nd_best, ...]
    Example: [3, 1, 5, 2, 4]
    """
    
    response = claude_api.generate(prompt, max_tokens=50)
    llm_ranking = parse_ranking(response)
    
    # Return best action according to LLM
    return top_actions[llm_ranking[0]]

This is the BIGGEST single improvement!


Day 19: Retrieval-Augmented Generation

RAG for similar examples:

def retrieve_similar_states(current_state, dataset, k=10):
    """Find k most similar states with successful actions."""
    
    # Embed all states
    state_embeddings = embed_all_states(dataset)
    current_emb = embed_state(current_state)
    
    # Cosine similarity
    similarities = cosine_similarity(current_emb, state_embeddings)
    top_k_idx = torch.topk(similarities, k).indices
    
    # Return successful examples
    examples = [
        dataset[i] for i in top_k_idx 
        if dataset[i].reward.terminal_success
    ]
    
    return examples

# Use in LLM prompt
similar = retrieve_similar_states(state, dataset, k=5)
prompt = f"""
Current state: {state}
Similar successful examples:
{format_examples(similar)}

Based on these, rank the candidate actions.
"""

Day 20: Chain-of-Thought Reasoning

Make LLM explain step-by-step:

prompt = f"""
State: {state}
Goal: {instruction}
Actions: {actions}

For each action, reason step-by-step:

Action 1: {action_1}
Step 1: What will happen physically?
Step 2: Will it achieve the goal?
Step 3: Any risks or failures?
Step 4: Overall rating (1-10):

[Repeat for all actions]

Final ranking: [best to worst]
"""

More expensive but more accurate


Day 21: Full System Evaluation

Test complete pipeline:

def evaluate_full_pipeline(dataset):
    results = 
    
    # 1. Baseline Transformer (no improvements)
    results['baseline'] = evaluate(transformer_basic)
    
    # 2. + Language
    results['language'] = evaluate(transformer_lang)
    
    # 3. + Data augmentation
    results['data_aug'] = evaluate(transformer_lang_aug)
    
    # 4. + Architecture improvements
    results['architecture'] = evaluate(multiscale_model)
    
    # 5. + Training improvements
    results['training'] = evaluate(trained_with_curriculum)
    
    # 6. + Ensemble
    results['ensemble'] = evaluate(ensemble_model)
    
    # 7. + LLM judge (FINAL)
    results['final'] = evaluate(system_with_llm_judge)
    
    return results

Expected final result: 60-70%+


πŸ“Š EXPECTED PROGRESS TRACKING

Checkpoint Selected Success Improvement Cumulative
Current Transformer 42-44% - Baseline
+Language (Day 4) 50-55% +8-11% +8-11%
+Data Aug (Day 7) 52-57% +2-5% +10-15%
+Architecture (Day 10) 54-59% +2-4% +12-17%
+Training (Day 14) 57-62% +3-5% +15-20%
+Ensemble (Day 16) 60-65% +3-5% +18-23%
+LLM Judge (Day 18) 65-75% +10-15% +23-33%
+RAG+CoT (Day 20) 67-78% +2-5% +25-36%

Final target: 65-75% selected success


πŸ’° API Cost Estimation

With unlimited API:

  • Embeddings: sentence-transformers (free, local)
  • Synthetic data: ~10K LLM calls
  • Explanations: ~56K LLM calls
  • LLM judge: ~3.5K calls/eval Γ— 10 evals = 35K calls
  • RAG: ~3.5K calls
  • CoT: ~3.5K calls (expensive, 500 tokens/call)

Total: ~110K LLM API calls over 3 weeks

With Claude API: ~$550-1,100 (at $5-10 per 1M tokens) Your case: Unlimited β†’ FREE! πŸŽ‰


🎯 SUCCESS CRITERIA

Minimum success (Week 2):

  • 55%+ selected success
  • Better than baseline (+12%)
  • Publishable improvement

Target (Week 3):

  • 60%+ selected success
  • Strong CVPR paper
  • Clear ablation study

Stretch (if LLM judge works well):

  • 70%+ selected success
  • SOTA-competitive at small scale
  • Major contribution

πŸ“‹ NEXT IMMEDIATE ACTIONS

Now (while current Transformer trains):

  1. βœ… Setup environment (pip install dependencies)
  2. βœ… Test language embedding generation
  3. βœ… Create implementation skeleton

When current training finishes (2h):

  1. Evaluate baseline (42-44%)
  2. Start Week 1 Day 1 (language integration)
  3. Launch parallel experiments

Ready to start implementation? πŸš€

Let me know when to begin Day 1, or I can start preparing now!