Precision Starts with Every Perception and Decision
Driven by a proprietary dual-system hierarchical architecture and the VTLA foundation model, Lumi-Mind utilizes multimodal inputs, multi-stage data alignment, and progressive training to empower hands with fine multi-finger dexterity, defining a new paradigm for dexterous manipulation.
def get_action(self, fused_tokens: torch.Tensor, state: torch.Tensor = None, embodiment_id: torch.LongTensor = None, action_mask: torch.Tensor = None):
# print(f"action_mask shape: {action_mask.shape if action_mask is not None else 'None'}")
# print(f"one sample action_mask: {action_mask[0] if action_mask is not None else 'None'}")
B = fused_tokens.size(0)
device = fused_tokens.device
if embodiment_id is None:
embodiment_id = torch.zeros(B, dtype=torch.long, device=device)
context_tokens = fused_tokens
if state is not None and self.state_encoder is not None:
state_emb = self.state_encoder(state, embodiment_id).unsqueeze(1)
context_tokens = torch.cat([context_tokens, state_emb], dim=1)
action_dim_total = getattr(self.config, "action_dim", None)
if action_dim_total is None:
action_dim_total = self.action_dim
if self.horizon > 1:
per_action_dim = getattr(self.config, "per_action_dim", action_dim_total // self.horizon)
else:
per_action_dim = action_dim_total
action = (torch.rand(B, action_dim_total, device=device) * 2 - 1)
# print(f"action shape: {action.shape}")
# print(f"one sample action: {action[0]}")
if self.horizon > 1:
action_seq = action.view(B, self.horizon, per_action_dim)
else:
action_seq = action.view(B, 1, per_action_dim)
action_mask = action_mask.view(B, 1, per_action_dim).repeat(1,self.horizon,1)
# print(f"action_mask: {action_mask}")
# print(f"one sample action_mask: {action_mask[0]}")
if action_mask is not None:
action_mask = action_mask.to(dtype=action_seq.dtype, device=action_seq.device)
assert action_mask.shape == action_seq.shape, f"action_mask shape {action_mask.shape} != noise shape {action_seq.shape}"
action_seq = action_seq * action_mask
else:
raise ValueError("action_mask must be provided for inference with flow matching.")
# print(f"action shape: {action_seq.shape}")
# print(f"one sample action: {action_seq[0]}")
N = int(getattr(self.config, "num_inference_timesteps", 32))
dt = 1.0 / N
for i in range(N):
t = i / N
time_index = int(t * 1000)
time_emb = self.time_pos_enc(1000)[:, time_index, :].to(device).squeeze(0)
time_emb = time_emb.unsqueeze(0).repeat(B, 1)
if self.horizon > 1 and self.action_encoder is not None:
action_seq = action_seq * action_mask
action_tokens = self.action_encoder(action_seq, embodiment_id)
else:
if hasattr(self, "single_action_proj"):
action_tokens = self.single_action_proj(action_seq)
else:
self.single_action_proj = nn.Linear(per_action_dim, self.embed_dim).to(device)
action_tokens = self.single_action_proj(action_seq)
x = action_tokens
for block in self.transformer_blocks:
x = block(x, context_tokens, time_emb)
x = self.norm_out(x)
if self.horizon > 1:
x_flat = x.reshape(B, -1)
if hasattr(self, "seq_pool_proj"):
x_pooled = self.seq_pool_proj(x_flat)
else:
self.seq_pool_proj = nn.Linear(self.horizon * self.embed_dim, self.embed_dim).to(device)
x_pooled = self.seq_pool_proj(x_flat)
else:
x_pooled = x.squeeze(1)
pred = self.mlp_head(x_pooled, embodiment_id)
action = action + dt * pred
if self.horizon > 1:
action_seq = action.view(B, self.horizon, per_action_dim)
else:
action_seq = action.view(B, 1, per_action_dim)
return action
