Views
No views yet
[x, y, yaw] waypoints for embodied person following.
 |
 |
 |
| Outdoor Obstacle-aware Tracking | Elevator Tracking | Underground Parking Tracking |
transformers>=4.56,<5.1from __future__ import annotations
2
3from pathlib import Path
4
5import torch
6from transformers import AutoModel, AutoTokenizer
7
8
9VISION_FEATURE_DIM = 1536
10HISTORY_FRAMES = 31
11COARSE_TOKENS_PER_FRAME = 4
12FINE_TOKENS_CURRENT_FRAME = 64
13
14
15class MiniCPMRobotTrackInference:
16 """Tokenizer and model wrapper for MiniCPM-RobotTrack inference."""
17
18 def __init__(
19 self,
20 checkpoint_path: str | Path = "openbmb/MiniCPM-RobotTrack",
21 device: str | torch.device | None = None,
22 ):
23 if device is None:
24 device = "cuda" if torch.cuda.is_available() else "cpu"
25 self.device = torch.device(device)
26 checkpoint = str(checkpoint_path)
27
28 self.tokenizer = AutoTokenizer.from_pretrained(checkpoint)
29 if self.tokenizer.pad_token_id is None:
30 self.tokenizer.pad_token = self.tokenizer.eos_token
31
32 self.model = AutoModel.from_pretrained(
33 checkpoint,
34 trust_remote_code=True,
35 )
36 self.model.to(self.device).eval()
37
38 @staticmethod
39 def _prepare_visual_tokens(
40 tokens: torch.Tensor,
41 time_indices: torch.Tensor,
42 name: str,
43 ) -> tuple[torch.Tensor, torch.Tensor]:
44 tokens = torch.as_tensor(tokens, dtype=torch.float32)
45 time_indices = torch.as_tensor(time_indices, dtype=torch.long)
46
47 if tokens.ndim == 2:
48 tokens = tokens.unsqueeze(0)
49 if time_indices.ndim == 1:
50 time_indices = time_indices.unsqueeze(0)
51 if tokens.ndim != 3 or tokens.shape[-1] != VISION_FEATURE_DIM:
52 raise ValueError(
53 f"{name}_tokens must have shape [B, N, {VISION_FEATURE_DIM}]"
54 )
55 if time_indices.shape != tokens.shape[:2]:
56 raise ValueError(
57 f"{name}_time_indices must match the first two dimensions of "
58 f"{name}_tokens"
59 )
60 return tokens, time_indices
61
62 @torch.inference_mode()
63 def predict(
64 self,
65 instruction: str,
66 coarse_tokens: torch.Tensor,
67 coarse_time_indices: torch.Tensor,
68 fine_tokens: torch.Tensor,
69 fine_time_indices: torch.Tensor,
70 ) -> torch.Tensor:
71 """Return eight ``[x, y, yaw]`` waypoints for each batch item."""
72 coarse_tokens, coarse_time_indices = self._prepare_visual_tokens(
73 coarse_tokens,
74 coarse_time_indices,
75 "coarse",
76 )
77 fine_tokens, fine_time_indices = self._prepare_visual_tokens(
78 fine_tokens,
79 fine_time_indices,
80 "fine",
81 )
82
83 batch_size = coarse_tokens.shape[0]
84 if fine_tokens.shape[0] != batch_size:
85 raise ValueError("coarse and fine feature batch sizes must match")
86
87 text = self.tokenizer(
88 [instruction] * batch_size,
89 return_tensors="pt",
90 padding=True,
91 truncation=True,
92 max_length=self.model.config.max_text_tokens,
93 )
94 outputs = self.model(
95 input_ids=text.input_ids.to(self.device),
96 attention_mask=text.attention_mask.to(self.device),
97 coarse_tokens=coarse_tokens.to(self.device),
98 coarse_time_indices=coarse_time_indices.to(self.device),
99 fine_tokens=fine_tokens.to(self.device),
100 fine_time_indices=fine_time_indices.to(self.device),
101 )
102 return outputs.trajectories.float().cpu()
103
104
105if __name__ == "__main__":
106 infer_runner = MiniCPMRobotTrackInference()
107
108 # Replace these placeholders with fused DINOv3 + SigLIP features produced
109 # by the project preprocessing pipeline.
110 coarse_tokens = torch.zeros(
111 HISTORY_FRAMES * COARSE_TOKENS_PER_FRAME,
112 VISION_FEATURE_DIM,
113 )
114 coarse_time_indices = torch.arange(HISTORY_FRAMES).repeat_interleave(
115 COARSE_TOKENS_PER_FRAME
116 )
117 fine_tokens = torch.zeros(
118 FINE_TOKENS_CURRENT_FRAME,
119 VISION_FEATURE_DIM,
120 )
121 fine_time_indices = torch.full(
122 (FINE_TOKENS_CURRENT_FRAME,),
123 HISTORY_FRAMES,
124 dtype=torch.long,
125 )
126
127 trajectory = infer_runner.predict(
128 instruction="Follow the person in the red shirt.",
129 coarse_tokens=coarse_tokens,
130 coarse_time_indices=coarse_time_indices,
131 fine_tokens=fine_tokens,
132 fine_time_indices=fine_time_indices,
133 )
134 print(trajectory) # [1, 8, 3]
135