🎮 仿真平台与Sim2Real¶

import mujoco
import mujoco.viewer
import numpy as np

# 加载机器人模型（MJCF格式）
model = mujoco.MjModel.from_xml_path("franka_emika_panda/scene.xml")
data = mujoco.MjData(model)

# 查看模型信息
print(f"广义坐标维度 (nq): {model.nq}")   # 关节位置维度
print(f"广义速度维度 (nv): {model.nv}")   # 关节速度维度
print(f"控制输入维度 (nu): {model.nu}")   # 执行器数量
print(f"刚体数量: {model.nbody}")

# 关节名称
for i in range(model.njnt):
    name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_JOINT, i)
    print(f"  关节{i}: {name}, 范围: [{model.jnt_range[i, 0]:.2f}, "
          f"{model.jnt_range[i, 1]:.2f}]")

# 仿真循环
with mujoco.viewer.launch_passive(model, data) as viewer:
    for step in range(10000):
        # 设置控制输入（关节力矩）
        data.ctrl[:7] = np.zeros(7)  # 7个关节力矩
        data.ctrl[7] = 0.04          # 夹爪开度

        # 前进一个仿真步 (默认dt=0.002s)
        mujoco.mj_step(model, data)

        # 读取状态
        joint_positions = data.qpos[:7].copy()
        joint_velocities = data.qvel[:7].copy()

        # 读取末端位姿
        ee_site_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, "end_effector")
        ee_pos = data.site_xpos[ee_site_id].copy()

        viewer.sync()

import gymnasium as gym
from gymnasium import spaces

class MuJoCoReachEnv(gym.Env):
    """MuJoCo机械臂到达任务（Gymnasium接口）"""

    metadata = {'render_modes': ['human', 'rgb_array']}

    def __init__(self, xml_path, render_mode=None):
        super().__init__()  # super()调用父类方法
        self.model = mujoco.MjModel.from_xml_path(xml_path)
        self.data = mujoco.MjData(self.model)
        self.render_mode = render_mode

        # 动作空间: 7个关节力矩
        self.action_space = spaces.Box(
            low=-1.0, high=1.0, shape=(7,), dtype=np.float32
        )

        # 观察空间: 关节角度(7) + 速度(7) + 末端位置(3) + 目标位置(3)
        self.observation_space = spaces.Box(
            low=-np.inf, high=np.inf, shape=(20,), dtype=np.float32
        )

        self.max_steps = 200
        self.current_step = 0
        self.target_pos = np.array([0.5, 0.0, 0.5])  # np.array创建NumPy数组

    def reset(self, seed=None, options=None):
        super().reset(seed=seed)

        mujoco.mj_resetData(self.model, self.data)

        # 随机化目标位置
        self.target_pos = np.array([
            self.np_random.uniform(0.3, 0.7),
            self.np_random.uniform(-0.3, 0.3),
            self.np_random.uniform(0.2, 0.6)
        ])

        self.current_step = 0
        return self._get_obs(), {}

    def step(self, action):
        # 缩放动作到实际力矩范围
        torque_limits = np.array([87, 87, 87, 87, 12, 12, 12])
        self.data.ctrl[:7] = action * torque_limits

        # 仿真多步（action repeat = 10, 即控制频率 = 50Hz）
        for _ in range(10):
            mujoco.mj_step(self.model, self.data)

        self.current_step += 1
        obs = self._get_obs()

        # 计算奖励
        ee_pos = self._get_ee_pos()
        dist = np.linalg.norm(ee_pos - self.target_pos)  # np.linalg线性代数运算
        reward = -dist  # 密集奖励

        if dist < 0.05:  # 到达目标
            reward += 10.0

        terminated = dist < 0.02
        truncated = self.current_step >= self.max_steps

        return obs, reward, terminated, truncated, {'distance': dist}

    def _get_obs(self):
        return np.concatenate([
            self.data.qpos[:7],          # 关节角度
            self.data.qvel[:7],          # 关节速度
            self._get_ee_pos(),          # 末端位置
            self.target_pos              # 目标位置
        ]).astype(np.float32)

    def _get_ee_pos(self):
        site_id = mujoco.mj_name2id(
            self.model, mujoco.mjtObj.mjOBJ_SITE, "end_effector"
        )
        return self.data.site_xpos[site_id].copy()

# Isaac Lab 核心优势: GPU上同时跑数千个仿真环境
# pip install isaaclab

import torch
from isaaclab.envs import ManagerBasedRLEnv
from isaaclab.managers import SceneEntityCfg

class FrankaReachEnvCfg:
    """Isaac Lab配置（声明式定义）"""

    # 场景配置
    scene = {
        "robot": {
            "type": "Articulation",
            "prim_path": "/World/envs/env_.*/Robot",
            "usd_path": "franka_panda.usd",
            "actuator_groups": {
                "panda_arm": {
                    "joint_names": ["panda_joint[1-7]"],
                    "effort_limit": 87.0,
                    "stiffness": 400.0,
                    "damping": 80.0
                }
            }
        },
        "target": {
            "type": "RigidObject",
            "prim_path": "/World/envs/env_.*/Target",
            "usd_path": "target_sphere.usd"
        }
    }

    # 并行环境数量
    num_envs = 4096

    # 仿真参数
    sim = {
        "dt": 1.0 / 120,           # 仿真频率120Hz
        "decimation": 2,            # 控制频率60Hz
        "physx": {
            "gpu_found_lost_pairs_capacity": 1024 * 1024,
            "gpu_total_aggregate_pairs_capacity": 16 * 1024
        }
    }

# 训练流水线（概念代码）
# env = ManagerBasedRLEnv(cfg=FrankaReachEnvCfg())
# obs = env.reset()  # (4096, obs_dim) — 4096个环境同时初始化！
#
# for step in range(total_steps):
#     actions = policy(obs)           # (4096, 7) — 批量推理
#     obs, rewards, dones, infos = env.step(actions)  # GPU上并行执行
#     # 在4096环境上同时收集经验 → 训练速度提升1000x+

# Isaac Sim/Lab 支持光线追踪渲染
# → 生成逼真的RGB/深度/分割图像 → 用于训练视觉策略

class CameraConfig:
    """仿真摄像头配置"""

    cameras = {
        "wrist_camera": {
            "prim_path": "/World/envs/env_.*/Robot/panda_hand/Camera",
            "resolution": (256, 256),
            "data_types": ["rgb", "depth", "semantic_segmentation"],
            "update_period": 0.05  # 20Hz
        },
        "overhead_camera": {
            "prim_path": "/World/envs/env_.*/OverheadCam",
            "resolution": (480, 640),
            "data_types": ["rgb", "depth"],
            "update_period": 0.1  # 10Hz
        }
    }

    # 在Isaac Sim中，可以启用RTX光追渲染
    # → 阴影、反射、透明物体都能正确渲染
    # → 对视觉Sim2Real至关重要

import numpy as np
from dataclasses import dataclass, field
from typing import Tuple

@dataclass  # @dataclass自动生成__init__等方法
class DomainRandomizationConfig:
    """Domain Randomization配置"""

    # 物理参数随机化
    friction_range: Tuple[float, float] = (0.3, 1.2)
    mass_scale_range: Tuple[float, float] = (0.5, 2.0)
    damping_scale_range: Tuple[float, float] = (0.8, 1.2)

    # 视觉随机化
    light_intensity_range: Tuple[float, float] = (0.3, 2.0)
    light_color_var: float = 0.2
    camera_pos_noise: float = 0.02      # ±2cm
    camera_rot_noise: float = 0.05      # ±2.8°
    texture_randomize: bool = True

    # 执行器随机化
    action_noise_std: float = 0.02
    action_delay_range: Tuple[int, int] = (0, 3)  # 延迟0-3步

    # 物体随机化
    object_size_scale: Tuple[float, float] = (0.8, 1.2)
    object_pos_noise: float = 0.05      # ±5cm

class DomainRandomizer:
    """Domain Randomization管理器"""

    def __init__(self, config: DomainRandomizationConfig):
        self.cfg = config

    def randomize_physics(self, model):
        """随机化物理参数"""
        # 摩擦系数
        for geom_id in range(model.ngeom):
            model.geom_friction[geom_id, 0] = np.random.uniform(
                *self.cfg.friction_range
            )

        # 物体质量
        for body_id in range(model.nbody):
            original_mass = model.body_mass[body_id]
            scale = np.random.uniform(*self.cfg.mass_scale_range)
            model.body_mass[body_id] = original_mass * scale

        # 关节阻尼
        for jnt_id in range(model.njnt):
            original_damp = model.dof_damping[jnt_id]
            scale = np.random.uniform(*self.cfg.damping_scale_range)
            model.dof_damping[jnt_id] = original_damp * scale

    def randomize_visual(self, renderer):
        """随机化视觉参数"""
        # 随机光照
        intensity = np.random.uniform(*self.cfg.light_intensity_range)
        color_noise = np.random.normal(0, self.cfg.light_color_var, 3)
        light_color = np.clip([1, 1, 1] + color_noise, 0, 1) * intensity
        renderer.set_light(color=light_color)

        # 随机相机扰动
        pos_noise = np.random.normal(0, self.cfg.camera_pos_noise, 3)
        rot_noise = np.random.normal(0, self.cfg.camera_rot_noise, 3)
        renderer.perturb_camera(pos_noise, rot_noise)

        # 随机纹理
        if self.cfg.texture_randomize:
            random_color = np.random.randint(0, 255, 3)
            renderer.set_table_color(random_color)

    def randomize_action(self, action):
        """随机化执行器输出（模拟真实执行器的不确定性）"""
        noisy_action = action + np.random.normal(
            0, self.cfg.action_noise_std, action.shape
        )
        return noisy_action

    def apply(self, env, episode_start=True):
        """在每个episode开始时应用随机化"""
        if episode_start:
            self.randomize_physics(env.model)
            self.randomize_visual(env.renderer)

            # 随机化物体初始位置
            for obj in env.objects:
                noise = np.random.uniform(
                    -self.cfg.object_pos_noise,
                    self.cfg.object_pos_noise, 3
                )
                obj.reset_position(obj.default_position + noise)

                # 随机化物体大小
                scale = np.random.uniform(*self.cfg.object_size_scale)
                obj.set_scale(scale)

import torch
import torch.nn as nn

class TeacherPolicy(nn.Module):  # 继承nn.Module定义网络层
    """Teacher策略（使用特权信息）"""

    def __init__(self, priv_obs_dim, act_dim, hidden=256):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(priv_obs_dim, hidden), nn.ELU(),
            nn.Linear(hidden, hidden), nn.ELU(),
            nn.Linear(hidden, hidden), nn.ELU(),
            nn.Linear(hidden, act_dim)
        )

    def forward(self, priv_obs):
        """
        priv_obs包含:
          - 关节角度/速度 (可观测)
          - 物体精确位姿 (特权，真实中需要视觉估计)
          - 接触力 (特权，真实中需要力矩传感器)
          - 摩擦系数 (特权，真实中不可知)
        """
        return self.net(priv_obs)

class StudentPolicy(nn.Module):
    """Student策略（只使用现实可观测信息）"""

    def __init__(self, obs_dim, act_dim, hidden=256, history_len=10):
        super().__init__()
        self.history_len = history_len

        # 观察编码器（处理历史观察）
        self.obs_encoder = nn.Sequential(
            nn.Linear(obs_dim * history_len, hidden), nn.ELU(),
            nn.Linear(hidden, hidden), nn.ELU()
        )

        # 动作输出
        self.action_head = nn.Sequential(
            nn.Linear(hidden, hidden), nn.ELU(),
            nn.Linear(hidden, act_dim)
        )

    def forward(self, obs_history):
        """
        obs_history: (B, history_len, obs_dim)
        只包含真实世界可获得的信息:
          - 本体感知 (关节编码器)
          - 历史动作
          - IMU数据
        """
        B = obs_history.shape[0]
        flat = obs_history.view(B, -1)  # 重塑张量形状
        features = self.obs_encoder(flat)
        return self.action_head(features)

class TeacherStudentTrainer:
    """Teacher-Student蒸馏训练"""

    def __init__(self, teacher, student, lr=1e-3):
        self.teacher = teacher
        self.student = student
        self.optimizer = torch.optim.Adam(student.parameters(), lr=lr)

    def distill_epoch(self, env, num_episodes=100):
        """
        1. 用Teacher在仿真中执行
        2. 收集 (student_obs, teacher_action) 数据对
        3. 训练Student模仿Teacher
        """
        dataset = []

        # 收集数据
        self.teacher.eval()
        for ep in range(num_episodes):
            priv_obs, student_obs = env.reset()
            obs_buffer = []

            for step in range(env.max_steps):
                # Teacher决策
                with torch.no_grad():  # 禁用梯度计算，节省内存
                    teacher_action = self.teacher(
                        torch.FloatTensor(priv_obs).unsqueeze(0)  # unsqueeze增加一个维度
                    ).squeeze(0).numpy()  # squeeze压缩维度

                # 记录Student视角的观察和Teacher的动作
                obs_buffer.append(student_obs)
                if len(obs_buffer) >= self.student.history_len:
                    history = np.stack(
                        obs_buffer[-self.student.history_len:]
                    )
                    dataset.append((history, teacher_action))

                # 执行动作
                priv_obs, student_obs, reward, done, info = \
                    env.step(teacher_action)

                if done:
                    break

        # 训练Student
        self.student.train()
        losses = []

        # 随机打乱
        np.random.shuffle(dataset)

        for i in range(0, len(dataset), 64):  # batch_size=64
            batch = dataset[i:i+64]
            obs_batch = torch.FloatTensor(np.stack([d[0] for d in batch]))
            act_batch = torch.FloatTensor(np.stack([d[1] for d in batch]))

            pred_actions = self.student(obs_batch)
            loss = nn.functional.mse_loss(pred_actions, act_batch)

            self.optimizer.zero_grad()  # 清零梯度
            loss.backward()  # 反向传播计算梯度
            self.optimizer.step()  # 更新参数
            losses.append(loss.item())  # 将单元素张量转为Python数值

        return np.mean(losses)

# 另一种Sim2Real方法: 让仿真尽可能接近真实

class SystemIdentification:
    """通过优化使仿真参数匹配真实数据"""

    def __init__(self, sim_env, real_trajectories):
        self.sim = sim_env
        self.real_data = real_trajectories

    def optimize_params(self, param_names, bounds, n_iter=100):
        """
        贝叶斯优化: 调整仿真参数使轨迹匹配真实数据
        """
        from scipy.optimize import differential_evolution

        def objective(params):
            # 设置仿真参数
            for name, value in zip(param_names, params):  # zip按位置配对
                self.sim.set_param(name, value)

            # 回放真实动作序列在仿真中
            total_error = 0
            for traj in self.real_data:
                self.sim.reset(traj['initial_state'])
                for action, real_next_state in zip(
                    traj['actions'], traj['next_states']
                ):
                    sim_next_state = self.sim.step(action)
                    total_error += np.linalg.norm(
                        sim_next_state - real_next_state
                    ) ** 2

            return total_error

        result = differential_evolution(objective, bounds, maxiter=n_iter)
        return dict(zip(param_names, result.x))