4. 环境搭建：基于PyTorch的仿真环境构建（MuJoCo/Gym接口）

说实话，搞强化学习控制，最头疼的往往不是算法本身，而是环境搭不起来。我记得刚入行那会儿，光配MuJoCo就折腾了两天，最后发现是Python版本不对。这节课，我把这些年踩过的坑和总结出来的套路，一次性说清楚。

4.1 为什么选MuJoCo + Gym？

你想想看，做力矩控制，我们需要什么？一个物理引擎，能算关节力矩、碰撞、摩擦。MuJoCo在这方面是行业标准。Gym呢，它把环境接口统一了，reset、step、render，一套API走天下。

我个人习惯用MuJoCo 2.1.0以上版本，配合Gymnasium（Gym的继承者）。为什么？因为老版Gym已经停止维护了，新项目直接用Gymnasium更省心。

4.2 环境搭建三步走

说白了，搭建过程就三步：装MuJoCo、装Gymnasium、写一个测试脚本。别想复杂了。

4.2.1 安装MuJoCo

MuJoCo从2.1.0开始开源了，直接pip就能装。我个人建议用conda创建虚拟环境，避免包冲突。

# 创建虚拟环境
conda create -n rl_control python=3.9
conda activate rl_control

# 安装MuJoCo
pip install mujoco

# 验证安装
python -c "import mujoco; print(mujoco.__version__)"

如果看到版本号，说明装好了。嗯，这里要注意，MuJoCo 2.1.0以上版本不再需要单独的license文件了，省了不少事。

4.2.2 安装Gymnasium和MuJoCo接口

Gymnasium官方提供了MuJoCo环境的接口包，叫gymnasium-robotics。不过我个人更推荐用gymnasium本身自带的MuJoCo环境，兼容性更好。

# 安装Gymnasium
pip install gymnasium

# 安装MuJoCo环境支持
pip install "gymnasium[mujoco]"

# 测试环境
python -c "import gymnasium as gym; env = gym.make('Ant-v4'); print(env.observation_space.shape)"

4.3 核心接口解析

Gymnasium的环境接口其实很简单，就几个函数。但每个函数的细节，决定了你的训练能不能跑通。

4.3.1 环境初始化与重置

每次训练开始，都要调用reset()。它会返回初始观测值。这里有个坑：reset()可以传seed参数，保证实验可重复。

import gymnasium as gym

env = gym.make('HalfCheetah-v4')
obs, info = env.reset(seed=42)  # 固定随机种子
print("观测维度:", obs.shape)
print("观测值:", obs[:10])  # 打印前10个维度

4.3.2 动作空间与观测空间

力矩控制的环境，动作空间通常是连续值，范围在[-1, 1]之间。观测空间则包含关节角度、角速度、质心位置等。

我个人习惯在训练前先打印一下空间信息，确认维度对不对。曾经有一次，我把动作维度搞错了，训练了三天才发现，损失函数一直不降。

4.3.3 step函数与奖励设计

step(action)是核心。它返回五个值：obs、reward、terminated、truncated、info。注意，Gymnasium把done拆成了terminated和truncated，前者是任务完成，后者是超时截断。

action = env.action_space.sample()  # 随机动作
obs, reward, terminated, truncated, info = env.step(action)

if terminated or truncated:
    print("回合结束，总奖励:", info.get('episode', {}).get('r', 0))

4.4 自定义MuJoCo环境

有时候，官方环境不够用。比如你想改个奖励函数，或者换个机器人模型。这时候就得自己写环境了。

4.4.1 继承gymnasium.Env

我建议直接继承gymnasium.Env，然后重写几个关键方法。下面是一个简单的自定义环境模板：

import gymnasium as gym
from gymnasium import spaces
import numpy as np
import mujoco

class CustomRobotEnv(gym.Env):
    def __init__(self, xml_path):
        super().__init__()
        self.model = mujoco.MjModel.from_xml_path(xml_path)
        self.data = mujoco.MjData(self.model)
        
        # 定义动作空间和观测空间
        self.action_space = spaces.Box(
            low=-1.0, high=1.0, 
            shape=(self.model.nu,), dtype=np.float32
        )
        self.observation_space = spaces.Box(
            low=-np.inf, high=np.inf,
            shape=(self.model.nq + self.model.nv,), dtype=np.float32
        )
    
    def reset(self, seed=None):
        super().reset(seed=seed)
        mujoco.mj_resetData(self.model, self.data)
        return self._get_obs(), {}
    
    def step(self, action):
        self.data.ctrl[:] = action
        mujoco.mj_step(self.model, self.data)
        obs = self._get_obs()
        reward = self._compute_reward()
        terminated = False
        truncated = False
        return obs, reward, terminated, truncated, {}
    
    def _get_obs(self):
        return np.concatenate([
            self.data.qpos.flat,
            self.data.qvel.flat
        ])
    
    def _compute_reward(self):
        # 自定义奖励函数
        return 0.0

4.5 与PyTorch的集成

环境搭好了，怎么跟PyTorch训练循环对接？其实很简单，就是把环境返回的numpy数组转成torch张量。

import torch
import gymnasium as gym

env = gym.make('HalfCheetah-v4')
obs, _ = env.reset()

# 转成PyTorch张量
obs_tensor = torch.from_numpy(obs).float().unsqueeze(0)

# 假设你有一个策略网络
action = policy_net(obs_tensor)
action_numpy = action.detach().numpy().squeeze()

# 执行动作
next_obs, reward, terminated, truncated, _ = env.step(action_numpy)

我个人习惯在训练循环里用torch.no_grad()包一下推理部分，省显存。另外，记得把reward也转成tensor，方便计算损失。

4.6 可视化与调试

训练的时候，光看loss曲线是不够的。你得看看机器人到底在干嘛。MuJoCo自带渲染器，但我觉得用Gymnasium的render模式更方便。

env = gym.make('Humanoid-v4', render_mode='human')
obs, _ = env.reset()

for _ in range(1000):
    action = env.action_space.sample()
    obs, reward, terminated, truncated, _ = env.step(action)
    if terminated or truncated:
        obs, _ = env.reset()

env.close()

4.7 本章知识体系

下面这张图，概括了本章的核心内容。从环境安装到自定义开发，再到与PyTorch的集成，每一步都有坑，但也都有解法。

嗯，环境搭建这部分，说白了就是「装好、测通、改对」。别怕出错，我当年第一次配环境，光编译就失败了七八次。但只要把基础打牢，后面的算法实现就会顺畅很多。