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