Github |
X |
Discord |
A Compact Vision-Language-Action Policy for Embodied target Tracking
MiniCPM-RobotTrack is a compact vision-language-action model built on MiniCPM4-0.5B for target tracking with the following highlights:[x, y, yaw] waypoints for embodied person following.
|
|
|
| Outdoor Obstacle-aware Tracking | Elevator Tracking | Underground Parking Tracking |
from __future__ import annotations
from pathlib import Path
import torch
from transformers import AutoModel, AutoTokenizer
VISION_FEATURE_DIM = 1536
HISTORY_FRAMES = 31
COARSE_TOKENS_PER_FRAME = 4
FINE_TOKENS_CURRENT_FRAME = 64
class MiniCPMRobotTrackInference:
"""Tokenizer and model wrapper for MiniCPM-RobotTrack inference."""
def __init__(
self,
checkpoint_path: str | Path = "openbmb/MiniCPM-RobotTrack",
device: str | torch.device | None = None,
):
if device is None:
device = "cuda" if torch.cuda.is_available() else "cpu"
self.device = torch.device(device)
checkpoint = str(checkpoint_path)
self.tokenizer = AutoTokenizer.from_pretrained(checkpoint)
if self.tokenizer.pad_token_id is None:
self.tokenizer.pad_token = self.tokenizer.eos_token
self.model = AutoModel.from_pretrained(
checkpoint,
trust_remote_code=True,
)
self.model.to(self.device).eval()
@staticmethod
def _prepare_visual_tokens(
tokens: torch.Tensor,
time_indices: torch.Tensor,
name: str,
) -> tuple[torch.Tensor, torch.Tensor]:
tokens = torch.as_tensor(tokens, dtype=torch.float32)
time_indices = torch.as_tensor(time_indices, dtype=torch.long)
if tokens.ndim == 2:
tokens = tokens.unsqueeze(0)
if time_indices.ndim == 1:
time_indices = time_indices.unsqueeze(0)
if tokens.ndim != 3 or tokens.shape[-1] != VISION_FEATURE_DIM:
raise ValueError(
f"{name}_tokens must have shape [B, N, {VISION_FEATURE_DIM}]"
)
if time_indices.shape != tokens.shape[:2]:
raise ValueError(
f"{name}_time_indices must match the first two dimensions of "
f"{name}_tokens"
)
return tokens, time_indices
@torch.inference_mode()
def predict(
self,
instruction: str,
coarse_tokens: torch.Tensor,
coarse_time_indices: torch.Tensor,
fine_tokens: torch.Tensor,
fine_time_indices: torch.Tensor,
) -> torch.Tensor:
"""Return eight ``[x, y, yaw]`` waypoints for each batch item."""
coarse_tokens, coarse_time_indices = self._prepare_visual_tokens(
coarse_tokens,
coarse_time_indices,
"coarse",
)
fine_tokens, fine_time_indices = self._prepare_visual_tokens(
fine_tokens,
fine_time_indices,
"fine",
)
batch_size = coarse_tokens.shape[0]
if fine_tokens.shape[0] != batch_size:
raise ValueError("coarse and fine feature batch sizes must match")
text = self.tokenizer(
[instruction] * batch_size,
return_tensors="pt",
padding=True,
truncation=True,
max_length=self.model.config.max_text_tokens,
)
outputs = self.model(
input_ids=text.input_ids.to(self.device),
attention_mask=text.attention_mask.to(self.device),
coarse_tokens=coarse_tokens.to(self.device),
coarse_time_indices=coarse_time_indices.to(self.device),
fine_tokens=fine_tokens.to(self.device),
fine_time_indices=fine_time_indices.to(self.device),
)
return outputs.trajectories.float().cpu()
if __name__ == "__main__":
infer_runner = MiniCPMRobotTrackInference()
# Replace these placeholders with fused DINOv3 + SigLIP features produced
# by the project preprocessing pipeline.
coarse_tokens = torch.zeros(
HISTORY_FRAMES * COARSE_TOKENS_PER_FRAME,
VISION_FEATURE_DIM,
)
coarse_time_indices = torch.arange(HISTORY_FRAMES).repeat_interleave(
COARSE_TOKENS_PER_FRAME
)
fine_tokens = torch.zeros(
FINE_TOKENS_CURRENT_FRAME,
VISION_FEATURE_DIM,
)
fine_time_indices = torch.full(
(FINE_TOKENS_CURRENT_FRAME,),
HISTORY_FRAMES,
dtype=torch.long,
)
trajectory = infer_runner.predict(
instruction="Follow the person in the red shirt.",
coarse_tokens=coarse_tokens,
coarse_time_indices=coarse_time_indices,
fine_tokens=fine_tokens,
fine_time_indices=fine_time_indices,
)
print(trajectory) # [1, 8, 3]
This project builds on and references MiniCPM, DINOv3, SigLIP, Habitat-Lab, Habitat-Sim, EVT-Bench, and TrackVLA. We thank the authors for their open-source contributions.
## License Model weights and code are open-sourced under the [Apache-2.0](./LICENSE) license.