用强化学习训练机器人行走

学习使用强化学习训练双足/四足机器人行走策略,涵盖环境搭建、奖励设计和Sim-to-Real迁移。

强化学习 机器人行走 双足机器人 四足机器人 Isaac Gym
用强化学习训练机器人行走

用强化学习训练机器人行走

让机器人稳定行走是机器人学的经典难题。传统方法需要复杂的动力学建模,而强化学习(RL)提供了一种数据驱动的解决方案。

为什么用强化学习?

传统行走控制 vs 强化学习:

方法优点缺点
传统控制可解释、稳定建模复杂、难以适应地形
强化学习自动学习、适应性强需要大量训练、黑盒

搭建训练环境

import numpy as np
import torch
from isaacgym import gymapi, gymtorch

class LeggedRobotEnv:
    def __init__(self, num_envs=4096, device="cuda"):
        self.num_envs = num_envs
        self.device = device

        # 创建Gym
        self.gym = gymapi.acquire_gym()

        # 仿真参数
        sim_params = gymapi.SimParams()
        sim_params.dt = 1/60
        sim_params.substeps = 2
        sim_params.up_axis = gymapi.UP_AXIS_Z
        sim_params.gravity = gymapi.Vec3(0, 0, -9.81)

        self.sim = self.gym.create_sim(
            0, 0, gymapi.SIM_PHYSX, sim_params
        )

        # 创建环境
        self._create_envs()

        # 获取张量
        self._acquire_tensors()

    def _create_envs(self):
        """创建并行环境"""
        # 添加地面
        plane_params = gymapi.PlaneParams()
        plane_params.normal = gymapi.Vec3(0, 0, 1)
        self.gym.add_ground(self.sim, plane_params)

        # 加载机器人资产(以Unitree A1为例)
        asset_root = "assets"
        asset_file = "urdf/a1/a1.urdf"
        asset_options = gymapi.AssetOptions()
        asset_options.fix_base_link = False
        asset_options.default_dof_drive_mode = gymapi.DOF_MODE_EFFORT

        robot_asset = self.gym.load_asset(
            self.sim, asset_root, asset_file, asset_options
        )

        # 创建环境网格
        spacing = 2.0
        env_lower = gymapi.Vec3(-spacing, -spacing, 0)
        env_upper = gymapi.Vec3(spacing, spacing, 0)

        self.envs = []
        self.actor_handles = []

        for i in range(self.num_envs):
            env = self.gym.create_env(self.sim, env_lower, env_upper, 64)

            # 放置机器人
            start_pose = gymapi.Transform()
            start_pose.p = gymapi.Vec3(0, 0, 0.5)

            actor_handle = self.gym.create_actor(
                env, robot_asset, start_pose, "robot", i, 1
            )
            self.envs.append(env)
            self.actor_handles.append(actor_handle)

        # 获取DOF属性
        self.num_dofs = self.gym.get_asset_dof_count(robot_asset)
        print(f"机器人自由度: {self.num_dofs}")

    def _acquire_tensors(self):
        """获取GPU张量"""
        self.root_state_tensor = self.gym.acquire_actor_root_state_tensor(self.sim)
        self.dof_state_tensor = self.gym.acquire_dof_state_tensor(self.sim)
        self.net_contact_force_tensor = self.gym.acquire_net_contact_force_tensor(self.sim)

        self.gym.refresh_actor_root_state_tensor(self.sim)
        self.gym.refresh_dof_state_tensor(self.sim)
        self.gym.refresh_net_contact_force_tensor(self.sim)

        # 包装为PyTorch张量
        self.root_states = gymtorch.wrap_tensor(self.root_state_tensor)
        self.dof_states = gymtorch.wrap_tensor(self.dof_state_tensor)
        self.contact_forces = gymtorch.wrap_tensor(self.net_contact_force_tensor)

    def reset(self):
        """重置环境"""
        # 随机化初始状态
        self.root_states[:, :3] = torch.tensor([0, 0, 0.5])
        self.root_states[:, 3:7] = torch.tensor([0, 0, 0, 1])  # 四元数
        self.root_states[:, 7:13] = 0  # 速度清零

        self.dof_states[:, 0] = 0  # 关节位置
        self.dof_states[:, 1] = 0  # 关节速度

        # 应用状态
        self.gym.set_actor_root_state_tensor(self.sim, self.root_state_tensor)
        self.gym.set_dof_state_tensor(self.sim, self.dof_state_tensor)

        self.gym.simulate(self.sim)
        self.gym.fetch_results(self.sim, True)

        return self._get_observations()

    def step(self, actions):
        """执行一步"""
        # 应用动作(关节力矩)
        self._apply_actions(actions)

        # 步进仿真
        self.gym.simulate(self.sim)
        self.gym.fetch_results(self.sim, True)

        # 刷新状态
        self.gym.refresh_actor_root_state_tensor(self.sim)
        self.gym.refresh_dof_state_tensor(self.sim)
        self.gym.refresh_net_contact_force_tensor(self.sim)

        # 获取观测、奖励、是否结束
        obs = self._get_observations()
        rewards = self._compute_rewards()
        dones = self._compute_dones()

        return obs, rewards, dones, {}

    def _get_observations(self):
        """获取状态观测"""
        # 基座姿态
        base_pos = self.root_states[:, :3]
        base_quat = self.root_states[:, 3:7]
        base_lin_vel = self.root_states[:, 7:10]
        base_ang_vel = self.root_states[:, 10:13]

        # 关节状态
        dof_pos = self.dof_states[:, 0].view(self.num_envs, -1)
        dof_vel = self.dof_states[:, 1].view(self.num_envs, -1)

        # 组合观测
        obs = torch.cat([
            base_pos,
            base_quat,
            base_lin_vel,
            base_ang_vel,
            dof_pos,
            dof_vel
        ], dim=-1)

        return obs

    def _compute_rewards(self):
        """计算奖励"""
        rewards = torch.zeros(self.num_envs, device=self.device)

        # 1. 前进奖励
        forward_vel = self.root_states[:, 7]  # X方向速度
        rewards += forward_vel * 1.0

        # 2. 存活奖励
        rewards += 0.1

        # 3. 能量惩罚
        energy = torch.sum(torch.abs(self.dof_states[:, 1]), dim=-1)
        rewards -= energy * 0.001

        # 4. 姿态惩罚(鼓励保持水平)
        base_quat = self.root_states[:, 3:7]
        # 计算倾斜角度
        tilt = torch.abs(base_quat[:, 0]) + torch.abs(base_quat[:, 1])
        rewards -= tilt * 0.1

        # 5. 跌倒惩罚
        base_height = self.root_states[:, 2]
        rewards += torch.where(base_height > 0.3, 0.0, -10.0)

        return rewards

    def _compute_dones(self):
        """计算是否结束"""
        # 跌倒检测
        base_height = self.root_states[:, 2]
        fallen = base_height < 0.2

        # 超时
        timeout = self.episode_length >= 1000

        return fallen | timeout

PPO训练

from stable_baselines3 import PPO
from stable_baselines3.common.vec_env import DummyVecEnv

# 创建环境
env = LeggedRobotEnv(num_envs=4096)

# 包装为SB3格式
vec_env = DummyVecEnv([lambda: env])

# 配置PPO
model = PPO(
    "MlpPolicy",
    vec_env,
    learning_rate=3e-4,
    n_steps=2048,
    batch_size=256,
    n_epochs=10,
    gamma=0.99,
    gae_lambda=0.95,
    clip_range=0.2,
    ent_coef=0.01,
    verbose=1
)

# 训练
model.learn(total_timesteps=10_000_000)

# 保存
model.save("walking_policy")

Sim-to-Real迁移

class WalkingPolicy:
    def __init__(self, model_path):
        self.model = PPO.load(model_path)

    def get_action(self, obs):
        """获取动作"""
        action, _ = self.model.predict(obs, deterministic=True)
        return action

    def deploy_on_real_robot(self, robot_interface):
        """部署到真实机器人"""
        obs = robot_interface.get_observation()
        done = False

        while not done:
            action = self.get_action(obs)
            robot_interface.send_command(action)

            # 等待控制周期
            time.sleep(0.02)  # 50Hz

            obs = robot_interface.get_observation()
            done = robot_interface.is_fallen()

# 使用
policy = WalkingPolicy("walking_policy")
policy.deploy_on_real_robot(real_robot)

FAQ

训练需要多长时间?

使用4096个并行环境,在RTX 3090上约需2-4小时。

如何提高行走稳定性?

增加地形随机化、调整奖励函数权重、使用课程学习。

可以训练跑步吗?

可以,但需要更高的控制频率和更精细的奖励设计。