具身智能仿真平台全景：Isaac Sim、MuJoCo、Habitat 技术对比与实战指南

"""
Isaac Sim 5.0 抓取放置任务示例
使用 Isaac Lab 框架进行强化学习训练
"""
from isaacsim import SimulationApp

# 启动仿真应用
sim = SimulationApp(
    launch_config={
        "headless": True,          # 无头模式（服务器训练）
        "multi_gpu": True,         # 多 GPU 并行
        "physics_device": "cuda",  # GPU 加速物理仿真
    }
)

from isaaclab.envs import VecEnvCfg, VecEnv
from isaaclab.tasks import ManipulationTask
from isaaclab.robots import FrankaArm, RobotiqGripper
from isaaclab.sensors import CameraSensor, TactileSensor
from isaaclab.utils import DomainRandomizer

# ── 1. 定义任务配置 ──
task_cfg = VecEnvCfg(
    num_envs=4096,           # 4096 个并行环境
    device="cuda:0",
    
    # 机器人配置
    robot=FrankaArm.with_gripper(RobotiqGripper()),
    
    # 任务定义：抓取红色方块放到蓝色区域
    task=ManipulationTask(
        object_type="cube",
        object_color="red",
        target_zone="blue_zone",
        success_threshold=0.02,  # 2cm 误差内算成功
    ),
    
    # 传感器配置
    sensors=[
        CameraSensor(
            name="wrist_cam",
            resolution=(128, 128),
            update_freq=10,       # 10Hz
            mount="robot_wrist",  # 安装在手腕上
        ),
        TactileSensor(
            name="finger_tactile",
            resolution=(32, 32),
            mount="gripper_fingers",
        ),
    ],
    
    # 奖励函数设计
    rewards={
        "approach_object": 1.0,    # 接近物体奖励
        "grasp_success": 10.0,     # 成功抓取奖励
        "reach_target": 50.0,      # 到达目标区域奖励
        "drop_penalty": -5.0,      # 掉落惩罚
        "collision_penalty": -1.0, # 碰撞惩罚
        "action_smoothness": 0.5,  # 动作平滑性奖励
    },
    
    #  episode 配置
    episode_length=200,  # 最大 200 步
    dt=1/60,             # 60Hz 仿真频率
)

# ── 2. 域随机化配置 ──
domain_rand = DomainRandomizer(
    randomize=[
        # 物体属性随机化
        {"target": "object_mass",    "range": (0.05, 0.5)},   # 50g-500g
        {"target": "object_friction", "range": (0.2, 0.9)},   # 摩擦系数
        {"target": "object_color",   "range": "hsv_full"},    # 全色相随机
        
        # 环境属性随机化
        {"target": "table_friction",  "range": (0.3, 0.8)},
        {"target": "lighting_intensity", "range": (200, 2000)},
        {"target": "camera_noise",    "range": (0.0, 0.05)},  # 传感器噪声
        
        # 机器人属性随机化
        {"target": "joint_damping",   "range": (0.9, 1.1)},   # ±10% 阻尼
        {"target": "gripper_force",   "range": (20, 60)},     # 抓取力
    ],
    resample_interval=20,  # 每 20 步重新随机化
)

# ── 3. 创建环境并训练 ──
env = VecEnv(task_cfg, domain_randomizer=domain_rand)

from isaaclab.agents import PPOAgent, PPOConfig

agent = PPOAgent(
    config=PPOConfig(
        lr=3e-4,
        gamma=0.99,
        clip_ratio=0.2,
        entropy_coef=0.01,
        batch_size=65536,        # 65K 步批量更新
        mini_batches=8,
        epochs_per_update=10,
        device="cuda:0",
    ),
    obs_space=env.observation_space,
    action_space=env.action_space,
)

# 训练循环
for iteration in range(1000):
    obs = env.reset()
    total_reward = 0
    
    for step in range(task_cfg.episode_length):
        action = agent.select_action(obs)
        next_obs, reward, done, info = env.step(action)
        agent.store_transition(obs, action, reward, done, next_obs)
        obs = next_obs
        total_reward += reward.mean()
    
    # PPO 更新
    metrics = agent.update()
    
    if iteration % 50 == 0:
        success_rate = info.get("success_rate", 0)
        print(f"Iter {iteration}: reward={total_reward:.1f}, "
              f"success={success_rate:.1%}")

# 保存策略
agent.save("franka_pick_place_v1.pt")
sim.close()

"""
MuJoCo 3.0 MJX 百万级并行仿真
在 GPU/TPU 上同时运行 100 万个仿真环境
"""
import jax
import jax.numpy as jnp
from mujoco import mjx
from mujoco.mjx import _src  # 底层 API

# ── 1. 加载 MuJoCo 模型 ──
import mujoco

# 从 XML 加载模型（例如：四足机器人）
mj_model = mujoco.MjModel.from_xml_path("mini_cheetah.xml")
mj_data = mujoco.MjData(mj_model)

# ── 2. 转换为 MJX 格式（JAX 后端）──
mjx_model = mjx.put_model(mj_model)
mjx_data = mjx.put_data(mj_model, mj_data)

# ── 3. 批量初始化 100 万个环境 ──
BATCH_SIZE = 1_000_000

# 使用 JAX vmap 批量初始化
def init_env(key):
    """初始化单个环境，带随机状态"""
    data = mjx_data.replace(
        qpos=jax.random.uniform(key, (mj_model.nq,)),
        qvel=jnp.zeros(mj_model.nv),
    )
    return data

keys = jax.random.split(jax.random.PRNGKey(0), BATCH_SIZE)
batch_data = jax.vmap(init_env)(keys)

# ── 4. 定义控制策略 ──
def policy(params, obs):
    """简单的线性策略（可被神经网络替代）"""
    # obs: [batch, obs_dim]
    # params: [obs_dim, action_dim]
    return obs @ params

# 初始化策略参数
obs_dim = mj_model.nq + mj_model.nv  # 位置 + 速度
action_dim = mj_model.nu              # 关节力矩
params = jax.random.normal(
    jax.random.PRNGKey(1), (obs_dim, action_dim)
) * 0.01

# ── 5. 批量仿真步骤 ──
@jax.jit
def sim_step(model, data, params):
    """单步仿真（JIT 编译，零 Python 开销）"""
    # 观测提取
    obs = jnp.concatenate([data.qpos, data.qvel])
    
    # 策略推理
    action = policy(params, obs)
    
    # 设置控制输入
    data = data.replace(ctrl=action)
    
    # 物理步进
    data = mjx.step(model, data)
    
    return data

# ── 6. 运行仿真 ──
import time

# JIT 编译（首次调用）
_ = sim_step(mjx_model, batch_data, params)

# 正式运行
start = time.time()
N_STEPS = 1000

data = batch_data
for _ in range(N_STEPS):
    data = sim_step(mjx_model, data, params)

# 同步等待完成
jax.block_until_ready(data)

elapsed = time.time() - start
env_steps = BATCH_SIZE * N_STEPS
sps = env_steps / elapsed

print(f"完成 {BATCH_SIZE:,} 个环境 × {N_STEPS} 步")
print(f"总计 {env_steps:,} 环境步")
print(f"速度: {sps:,.0f} 环境步/秒")
print(f"等效仿真时间: {elapsed * 0.002:.1f} 小时 "
      f"(假设 2ms/步的真实频率)")

"""
Sim-to-Real 流水线：MuJoCo 训练 → Isaac Sim 验证 → 真实部署
"""
import mujoco
from mujoco import mjx
import jax
import jax.numpy as jnp
import numpy as np

# ═══════════════════════════════════════════════════
# 阶段 1：MuJoCo 基础策略训练
# ═══════════════════════════════════════════════════

class SimToRealPipeline:
    def __init__(self, mjcf_path: str):
        # 加载 MuJoCo 模型
        self.mj_model = mujoco.MjModel.from_xml_path(mjcf_path)
        self.mjx_model = mjx.put_model(self.mj_model)
        
        # 配置并行环境
        self.batch_size = 500_000
        self.obs_dim = self.mj_model.nq + self.mj_model.nv
        self.act_dim = self.mj_model.nu
        
    def train_base_policy(self, n_iterations=5000):
        """阶段 1：在 MuJoCo 中训练基础策略"""
        from isaaclab.agents import PPOAgent
        
        # 初始化 PPO 智能体
        agent = PPOAgent(
            obs_dim=self.obs_dim,
            act_dim=self.act_dim,
            hidden_dims=[256, 256, 128],
            learning_rate=3e-4,
        )
        
        # 并行仿真环境
        key = jax.random.PRNGKey(42)
        env_keys = jax.random.split(key, self.batch_size)
        
        # 初始化环境数据
        batch_data = jax.vmap(self._init_env)(env_keys)
        
        for iteration in range(n_iterations):
            # 收集经验
            obs = self._extract_obs(batch_data)
            actions = agent.select_action(obs)
            
            # 批量仿真步进
            batch_data = jax.vmap(self._step)(
                jax.random.fold_in(key, iteration),
                batch_data, actions
            )
            
            # 存储并更新
            next_obs = self._extract_obs(batch_data)
            rewards = self._compute_rewards(batch_data)
            dones = self._compute_dones(batch_data)
            
            agent.store_batch(obs, actions, rewards, dones, next_obs)
            agent.update()
            
            # 重置已完成的环境
            batch_data = jax.vmap(self._reset_if_done)(
                env_keys, batch_data, dones
            )
            
            if iteration % 100 == 0:
                success = (rewards > 50).mean()
                print(f"Iter {iteration}: success_rate={success:.1%}")
        
        return agent
    
    def _init_env(self, key):
        """初始化单个环境"""
        data = mjx.put_data(self.mj_model, 
                           mujoco.MjData(self.mj_model))
        # 随机化物体位置
        obj_pos = jax.random.uniform(key, (3,), 
                                     min=jnp.array([-0.2, -0.1, 0.0]),
                                     max=jnp.array([0.2, 0.1, 0.0]))
        return data
    
    def _step(self, key, data, action):
        """单步仿真"""
        data = data.replace(ctrl=action)
        return mjx.step(self.mjx_model, data)
    
    def _extract_obs(self, data):
        """提取观测"""
        return jnp.concatenate([data.qpos, data.qvel], axis=-1)
    
    def _compute_rewards(self, data):
        """计算奖励（抓取任务）"""
        # 简化版：基于末端执行器到物体距离的奖励
        ee_pos = data.site_xpos[:, 0]  # 末端位置
        obj_pos = data.body_pos[:, 1]  # 物体位置
        dist = jnp.linalg.norm(ee_pos - obj_pos, axis=-1)
        return jnp.where(dist < 0.02, 100.0, -dist * 10)
    
    def _compute_dones(self, data):
        """计算终止条件"""
        return data.time > 10.0  # 10 秒超时
    
    def _reset_if_done(self, key, data, done):
        """条件重置"""
        init_data = self._init_env(key)
        return jax.tree.map(
            lambda new, old: jnp.where(done, new, old),
            init_data, data
        )


# ═══════════════════════════════════════════════════
# 阶段 2：Isaac Sim 视觉域随机化验证
# ═══════════════════════════════════════════════════

def validate_in_isaac_sim(policy_path: str):
    """阶段 2：在 Isaac Sim 中验证策略的视觉鲁棒性"""
    from isaacsim import SimulationApp
    from isaaclab.envs import VecEnv
    
    sim = SimulationApp({"headless": True, "multi_gpu": True})
    
    # 加载 MuJoCo 训练的策略
    from isaaclab.agents import PPOAgent
    agent = PPOAgent.load(policy_path)
    
    # 在 Isaac Sim 中创建相同任务
    env = VecEnv(
        task="franka_desktop_grasp",
        num_envs=256,
        # 视觉域随机化
        domain_randomization={
            "lighting": {"intensity": (200, 3000), "color_temp": (2700, 6500)},
            "object_texture": "random",
            "background": "random_indoor",
            "camera_noise": {"position": 0.01, "orientation": 0.5},
        },
        sensors=["rgb_camera", "depth_camera"],
    )
    
    # 评估策略
    total_success = 0
    n_episodes = 1000
    
    for ep in range(n_episodes):
        obs = env.reset()
        ep_reward = 0
        
        for step in range(200):
            action = agent.select_action(obs)
            obs, reward, done, info = env.step(action)
            ep_reward += reward.mean()
            
            if done.all():
                break
        
        if info.get("success", False):
            total_success += 1
    
    success_rate = total_success / n_episodes
    print(f"Isaac Sim 验证成功率: {success_rate:.1%}")
    
    sim.close()
    return success_rate


# ═══════════════════════════════════════════════════
# 阶段 3：真实机器人部署
# ═══════════════════════════════════════════════════

def deploy_to_real_robot(policy_path: str, robot_ip: str):
    """阶段 3：部署到真实 Franka 机器人"""
    import frankx  # Franka 机器人 Python SDK
    
    # 连接机器人
    robot = frankx.Robot(robot_ip)
    robot.set_default_behavior()
    robot.recover_from_errors()
    
    # 加载策略
    from isaaclab.agents import PPOAgent
    agent = PPOAgent.load(policy_path)
    
    # 安全限制
    MAX_FORCE = 30.0      # 关节力矩上限 (Nm)
    MAX_VELOCITY = 0.5    # 末端速度上限 (m/s)
    WORKSPACE_BOUNDS = {  # 工作空间限制
        "x": (-0.5, 0.5),
        "y": (-0.3, 0.3),
        "z": (0.0, 0.8),
    }
    
    print("🚀 开始真实机器人部署...")
    print("⚠️ 请确保工作空间内无人员")
    
    # 主控制循环
    state = robot.read_state()
    
    for step in range(500):  # 最多 500 步（约 8 秒）
        # 提取观测
        obs = np.concatenate([
            state.q,           # 关节位置
            state.dq,          # 关节速度
        ])
        
        # 策略推理
        action = agent.select_action(obs)
        
        # 安全检查
        action = np.clip(action, -MAX_FORCE, MAX_FORCE)
        
        # 执行动作
        robot.apply_torque(action)
        
        # 更新状态
        state = robot.read_state()
        
        # 检查是否到达目标
        ee_pos = state.O_T_EE[:3, 3]
        if np.linalg.norm(ee_pos - TARGET_POSITION) < 0.02:
            print("✅ 任务成功！")
            robot.gripper.grasp(0.04)  # 闭合夹爪
            break
    
    robot.stop()
    print("🏁 部署完成")


# ═══════════════════════════════════════════════════
# 主流程
# ═══════════════════════════════════════════════════

if __name__ == "__main__":
    # 阶段 1：MuJoCo 训练
    pipeline = SimToRealPipeline("franka_grasp.xml")
    agent = pipeline.train_base_policy(n_iterations=5000)
    agent.save("franka_grasp_mujoco.pt")
    
    # 阶段 2：Isaac Sim 验证
    success_rate = validate_in_isaac_sim("franka_grasp_mujoco.pt")
    
    if success_rate > 0.7:  # 成功率 > 70% 才部署
        # 阶段 3：真实部署
        deploy_to_real_robot("franka_grasp_mujoco.pt", "192.168.1.100")
    else:
        print(f"❌ 成功率 {success_rate:.1%} < 70%，暂不部署")