Views
No views yet

pip install torch transformers safetensors1import torch
2from transformers import AutoModel, AutoTokenizer
3
4# Load model
5model = AutoModel.from_pretrained(
6 "DolphinGR00T-N1.5-3B-Zero",
7 trust_remote_code=True,
8 torch_dtype="auto"
9)
10
11# Load tokenizer
12tokenizer = AutoTokenizer.from_pretrained("DolphinGR00T-N1.5-3B-Zero")
13
14# Move to GPU if available
15device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
16model = model.to(device)1import torch
2import torch.nn.functional as F
3from PIL import Image
4import numpy as np
5
6def prepare_image(image_path, target_size=(224, 224)):
7 """Prepare image for model input"""
8 image = Image.open(image_path).convert('RGB')
9 image = image.resize(target_size)
10 # Normalize to [-1, 1]
11 image = np.array(image).astype(np.float32) / 127.5 - 1.0
12 image = torch.from_numpy(image).permute(2, 0, 1)
13 return image
14
15def inference(model, tokenizer, image_paths, instruction, robot_state, device):
16 """
17 Run inference to generate robot actions
18
19 Args:
20 image_paths: List of paths to camera images
21 instruction: Natural language instruction
22 robot_state: Current robot proprioception (joint angles, etc.)
23 device: torch device
24
25 Returns:
26 actions: Predicted robot actions
27 """
28 model.eval()
29
30 with torch.no_grad():
31 # Prepare inputs
32 images = torch.stack([prepare_image(path) for path in image_paths])
33 images = images.unsqueeze(0).to(device) # Add batch dimension
34
35 # Tokenize instruction
36 text_inputs = tokenizer(
37 instruction,
38 return_tensors="pt",
39 padding=True,
40 truncation=True,
41 max_length=256
42 ).to(device)
43
44 # Robot state (example: 32-dim joint angles)
45 if isinstance(robot_state, list):
46 robot_state = torch.tensor(robot_state, dtype=torch.float32)
47 robot_state = robot_state.unsqueeze(0).to(device)
48
49 # Forward pass through backbone
50 # Note: This is a simplified example - actual implementation depends on model interface
51 vision_features = model.backbone.eagle_model.vision_model(images)
52
53 # Process language
54 language_features = model.backbone.eagle_model.language_model.model(
55 input_ids=text_inputs.input_ids,
56 attention_mask=text_inputs.attention_mask
57 ).last_hidden_state
58
59 # Combine features (simplified - actual fusion may be more complex)
60 combined_features = torch.cat([
61 vision_features.mean(dim=1), # Pool vision features
62 language_features.mean(dim=1) # Pool language features
63 ], dim=-1)
64
65 # Generate actions through diffusion process
66 # This is a simplified placeholder - actual diffusion requires multiple steps
67 action_features = model.action_head.model(
68 combined_features,
69 timesteps=torch.zeros(1, device=device),
70 context=robot_state
71 )
72
73 # Decode to action space
74 actions = model.action_head.action_decoder(action_features)
75
76 return actions
77
78# Example usage
79image_paths = ["camera1.jpg", "camera2.jpg"]
80instruction = "Pick up the red cube and place it on the table"
81robot_state = torch.randn(32) # Example: 32 joint angles
82
83actions = inference(model, tokenizer, image_paths, instruction, robot_state, device)
84print(f"Predicted actions shape: {actions.shape}")1import torch
2import torch.nn as nn
3from torch.utils.data import DataLoader, Dataset
4from transformers import get_linear_schedule_with_warmup
5
6class RobotDataset(Dataset):
7 """Example dataset for robot manipulation tasks"""
8 def __init__(self, data_path, tokenizer, transform=None):
9 self.data = [] # Load your data here
10 self.tokenizer = tokenizer
11 self.transform = transform
12
13 def __len__(self):
14 return len(self.data)
15
16 def __getitem__(self, idx):
17 # Return dict with keys: images, instruction, robot_state, target_actions
18 sample = self.data[idx]
19
20 # Process images
21 images = torch.stack([self.transform(img) for img in sample['images']])
22
23 # Tokenize instruction
24 text = self.tokenizer(
25 sample['instruction'],
26 return_tensors="pt",
27 padding="max_length",
28 truncation=True,
29 max_length=256
30 )
31
32 return {
33 'images': images,
34 'input_ids': text['input_ids'].squeeze(),
35 'attention_mask': text['attention_mask'].squeeze(),
36 'robot_state': torch.tensor(sample['robot_state'], dtype=torch.float32),
37 'target_actions': torch.tensor(sample['target_actions'], dtype=torch.float32)
38 }
39
40def train_step(model, batch, criterion, device):
41 """Single training step"""
42 # Move batch to device
43 images = batch['images'].to(device)
44 input_ids = batch['input_ids'].to(device)
45 attention_mask = batch['attention_mask'].to(device)
46 robot_state = batch['robot_state'].to(device)
47 target_actions = batch['target_actions'].to(device)
48
49 # Forward pass (simplified - actual implementation may differ)
50 # Process vision
51 vision_features = model.backbone.eagle_model.vision_model(images)
52
53 # Process language
54 language_output = model.backbone.eagle_model.language_model.model(
55 input_ids=input_ids,
56 attention_mask=attention_mask
57 )
58 language_features = language_output.last_hidden_state
59
60 # Combine modalities
61 combined_features = torch.cat([
62 vision_features.mean(dim=1),
63 language_features.mean(dim=1)
64 ], dim=-1)
65
66 # Generate actions (simplified diffusion)
67 predicted_actions = model.action_head(
68 combined_features,
69 context=robot_state
70 )
71
72 # Compute loss
73 loss = criterion(predicted_actions, target_actions)
74
75 return loss
76
77# Training setup
78def train_model(model, train_dataset, val_dataset, config):
79 """Main training loop"""
80 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
81 model = model.to(device)
82
83 # Create dataloaders
84 train_loader = DataLoader(
85 train_dataset,
86 batch_size=config['batch_size'],
87 shuffle=True,
88 num_workers=4
89 )
90
91 val_loader = DataLoader(
92 val_dataset,
93 batch_size=config['batch_size'],
94 shuffle=False,
95 num_workers=4
96 )
97
98 # Setup optimizer with different learning rates for backbone and action head
99 optimizer = torch.optim.AdamW([
100 {'params': model.backbone.parameters(), 'lr': config['backbone_lr']},
101 {'params': model.action_head.parameters(), 'lr': config['action_head_lr']}
102 ], weight_decay=config['weight_decay'])
103
104 # Learning rate scheduler
105 num_training_steps = len(train_loader) * config['num_epochs']
106 scheduler = get_linear_schedule_with_warmup(
107 optimizer,
108 num_warmup_steps=config['warmup_steps'],
109 num_training_steps=num_training_steps
110 )
111
112 # Loss function
113 criterion = nn.MSELoss() # or nn.L1Loss() for action prediction
114
115 # Training loop
116 for epoch in range(config['num_epochs']):
117 model.train()
118 total_loss = 0
119
120 for batch_idx, batch in enumerate(train_loader):
121 optimizer.zero_grad()
122
123 loss = train_step(model, batch, criterion, device)
124
125 loss.backward()
126
127 # Gradient clipping
128 torch.nn.utils.clip_grad_norm_(
129 model.parameters(),
130 config['max_grad_norm']
131 )
132
133 optimizer.step()
134 scheduler.step()
135
136 total_loss += loss.item()
137
138 if batch_idx % config['log_interval'] == 0:
139 print(f"Epoch {epoch}, Batch {batch_idx}, Loss: {loss.item():.4f}")
140
141 # Validation
142 model.eval()
143 val_loss = 0
144 with torch.no_grad():
145 for batch in val_loader:
146 loss = train_step(model, batch, criterion, device)
147 val_loss += loss.item()
148
149 avg_train_loss = total_loss / len(train_loader)
150 avg_val_loss = val_loss / len(val_loader)
151
152 print(f"Epoch {epoch}: Train Loss: {avg_train_loss:.4f}, Val Loss: {avg_val_loss:.4f}")
153
154 # Save checkpoint
155 if (epoch + 1) % config['save_interval'] == 0:
156 torch.save({
157 'epoch': epoch,
158 'model_state_dict': model.state_dict(),
159 'optimizer_state_dict': optimizer.state_dict(),
160 'scheduler_state_dict': scheduler.state_dict(),
161 'train_loss': avg_train_loss,
162 'val_loss': avg_val_loss,
163 }, f"checkpoint_epoch_{epoch+1}.pt")
164
165# Example configuration
166config = {
167 'batch_size': 16,
168 'num_epochs': 100,
169 'backbone_lr': 1e-5,
170 'action_head_lr': 1e-4,
171 'weight_decay': 0.01,
172 'warmup_steps': 1000,
173 'max_grad_norm': 1.0,
174 'log_interval': 10,
175 'save_interval': 10
176}
177
178# Create dataset (you need to implement data loading)
179# train_dataset = RobotDataset("path/to/train/data", tokenizer)
180# val_dataset = RobotDataset("path/to/val/data", tokenizer)
181
182# Train model
183# train_model(model, train_dataset, val_dataset, config)torch.cuda.amp for faster training:1from torch.cuda.amp import GradScaler, autocast
2
3scaler = GradScaler()
4
5with autocast():
6 loss = train_step(model, batch, criterion, device)
7
8scaler.scale(loss).backward()
9scaler.step(optimizer)
10scaler.update()model.backbone.eagle_model.language_model.gradient_checkpointing_enable()1# Freeze backbone
2for param in model.backbone.parameters():
3 param.requires_grad = False
4
5# Train only action head
6optimizer = torch.optim.AdamW(
7 model.action_head.parameters(),
8 lr=1e-4
9)1@software{DolphinGR00T2024,
2 title={DolphinGR00T-N1.5-3B-Zero: a Permissively Licensed Reimplementation of GR00T-N1.5-3B},
3 author={Eric Hartford},
4 year={2024},
5 license={Apache-2.0}
6}