从任意姿态到任意目标构型：智元 BFM-2 运动基座模型的技术解读

"""
简易条件轨迹生成器：给定 (q_start, q_goal)，生成 N 步过渡轨迹
模拟单臂 6-DOF 机器人，在关节空间做连续概率建模
"""
import torch
import torch.nn as nn
import numpy as np

# ---- 1. 生成训练数据：随机 start/goal + 插值轨迹 ----
def make_trajectory(q_start, q_goal, steps=50):
    """线性插值生成参考轨迹（真实数据应来自动力学仿真或示教）"""
    t = np.linspace(0, 1, steps)
    traj = q_start * (1 - t[:, None]) + q_goal * t[:, None]
    return traj  # shape: (steps, dof)

dof = 6
num_samples = 5000
steps = 50

# 随机采样关节角 [-π, π]
q_starts = np.random.uniform(-np.pi, np.pi, (num_samples, dof))
q_goals = np.random.uniform(-np.pi, np.pi, (num_samples, dof))

trajectories = np.stack([
    make_trajectory(q_starts[i], q_goals[i], steps)
    for i in range(num_samples)
])  # (num_samples, steps, dof)

# ---- 2. 条件轨迹生成网络 ----
class TrajectoryGenerator(nn.Module):
    """
    输入: q_start (dof) + q_goal (dof) → 编码为条件向量
    输出: steps × dof 的过渡轨迹
    """
    def __init__(self, dof=6, steps=50, hidden=128):
        super().__init__()
        self.dof = dof
        self.steps = steps
        # 条件编码器：把 start + goal 映射到 hidden 维特征
        self.cond_encoder = nn.Sequential(
            nn.Linear(dof * 2, hidden),
            nn.ReLU(),
            nn.Linear(hidden, hidden),
            nn.ReLU(),
        )
        # 轨迹解码器：从条件特征 + 时间步生成关节角
        self.decoder = nn.Sequential(
            nn.Linear(hidden + 1, hidden),  # +1 for time token
            nn.ReLU(),
            nn.Linear(hidden, dof),
        )

    def forward(self, q_start, q_goal):
        cond = self.cond_encoder(torch.cat([q_start, q_goal], dim=-1))
        # 为每个时间步生成输出
        t_tokens = torch.linspace(0, 1, self.steps, device=q_start.device)
        traj = []
        for t in t_tokens:
            t_input = torch.cat([cond, t.unsqueeze(0).expand(cond.shape[0], 1)], dim=-1)
            traj.append(self.decoder(t_input))
        return torch.stack(traj, dim=1)  # (batch, steps, dof)

# ---- 3. 训练 ----
model = TrajectoryGenerator(dof=dof, steps=steps)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

q_start_tensor = torch.tensor(q_starts, dtype=torch.float32)
q_goal_tensor = torch.tensor(q_goals, dtype=torch.float32)
traj_tensor = torch.tensor(trajectories, dtype=torch.float32)

for epoch in range(200):
    model.train()
    pred = model(q_start_tensor, q_goal_tensor)
    loss = nn.MSELoss()(pred, traj_tensor)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if epoch % 50 == 0:
        print(f"Epoch {epoch}: loss={loss.item():.4f}")

# ---- 4. 推理：从任意姿态到任意目标 ----
model.eval()
test_start = torch.tensor([[0.5, -1.2, 0.8, -0.3, 1.0, 0.2]], dtype=torch.float32)
test_goal = torch.tensor([[1.5, 0.5, -0.8, 0.6, -1.0, 1.5]], dtype=torch.float32)

generated_traj = model(test_start, test_goal).detach().numpy()
print(f"生成轨迹形状: {generated_traj.shape}")  # (1, 50, 6)
print(f"起点关节角: {generated_traj[0, 0]}")
print(f"终点关节角: {generated_traj[0, -1]}")

pip install torch numpy
python trajectory_generator.py