4. 策略梯度方法入门

任务目标

使用REINFORCE策略梯度算法,在简化四足机器人仿真环境中训练小跑步态策略。智能体通过关节角度观测与环境交互,策略网络直接输出关节目标角度,经奖励信号反馈优化策略参数,最终使机器人形成稳定的对角小跑步态(左前+右后、右前+左后交替触地)。

环境与依赖

  • Python 3.8+,仅需NumPy(仿真与训练均基于numpy实现,不依赖外部物理引擎)
  • 核心组件:SimpleQuadrupedEnv(简化四足动力学)、PolicyNetwork(高斯策略)、REINFORCE训练循环

实现原理

REINFORCE的核心是策略梯度定理:∇J(θ) = E[∇θ log π(a|s,θ) · R(τ)]。本案例中策略网络输出动作均值,方差固定;每回合结束后用累计折扣回报计算梯度并更新参数。奖励函数设计为:前进速度奖励 + 身体高度惩罚 + 动作平滑惩罚 + 小跑节奏奖励(对角腿相位差接近π/2时给予额外奖励)。

完整实现代码

<code>import numpy as np

# ---------- 简化四足机器人环境 ----------
class SimpleQuadrupedEnv:
    """8自由度(4腿×2关节:髋/膝),对角小跑目标"""
    def __init__(self):
        self.dt = 0.02
        self.action_dim = 8
        self.obs_dim = 12  # 8关节角 + 身体俯仰/横滚角速度 + 前进速度
        
        # 腿索引: 0=左前(LF), 1=右前(RF), 2=左后(LH), 3=右后(RH)
        self.diag_pairs = [(0,3), (1,2)]  # 对角配对: LF+RH, RF+LH
        
        self.reset()
    
    def reset(self):
        # 关节角初始为站立姿态(髋0.2rad, 膝-0.4rad)
        self.joint_angles = np.tile([0.2, -0.4], 4) + np.random.normal(0, 0.02, 8)
        self.body_pitch_rate = 0.0
        self.body_roll_rate = 0.0
        self.forward_vel = 0.0
        self.step_count = 0
        self.phase = 0.0  # 步态相位 [0, 2π)
        return self._get_obs()
    
    def _get_obs(self):
        return np.concatenate([
            self.joint_angles,
            [self.body_pitch_rate, self.body_roll_rate, self.forward_vel],
            [np.sin(self.phase), np.cos(self.phase)]
        ])
    
    def step(self, action):
        # 动作范围限制 [-1, 1] → 映射到关节角度变化
        target_angles = np.clip(action, -1, 1) * 0.6 + np.tile([0.2, -0.4], 4)
        
        # 简单动力学:关节向目标角度移动,产生地面反作用力
        angle_diff = target_angles - self.joint_angles
        self.joint_angles += angle_diff * 0.3
        
        # 计算腿的触地状态(简化:髋关节角度决定摆动/支撑)
        hip_angles = self.joint_angles[0::2]
        swing = hip_angles > 0.3  # 髋角大于0.3视为摆动相
        
        # 身体速度:支撑腿推动身体前进
        support_legs = ~swing
        forward_force = np.sum(support_legs.reshape(2,2).any(axis=1)) * 0.5
        self.forward_vel = self.forward_vel * 0.9 + forward_force * 0.1
        
        # 身体姿态扰动(简化)
        self.body_pitch_rate = np.random.normal(0, 0.01) + (np.sum(hip_angles[:2]) - np.sum(hip_angles[2:])) * 0.05
        self.body_roll_rate = np.random.normal(0, 0.01) + (hip_angles[0] - hip_angles[1]) * 0.05
        
        # 步态相位推进(目标频率2Hz)
        self.phase = (self.phase + 2 * np.pi * 2 * self.dt) % (2 * np.pi)
        self.step_count += 1
        
        # ---------- 奖励计算 ----------
        # 1. 前进速度奖励
        vel_reward = self.forward_vel * 2.0
        
        # 2. 身体姿态惩罚(保持水平)
        posture_penalty = -abs(self.body_pitch_rate) * 0.5 - abs(self.body_roll_rate) * 0.5
        
        # 3. 动作平滑惩罚
        if self.step_count > 1:
            smooth_penalty = -np.mean((action - self.last_action)**2) * 0.1
        else:
            smooth_penalty = 0.0
        self.last_action = action.copy()
        
        # 4. 小跑节奏奖励:对角腿相位差接近π
        # 计算当前对角腿是否处于相反相位(一个摆动一个支撑)
        lf_swing, rf_swing = swing[0], swing[1]
        lh_swing, rh_swing = swing[2], swing[3]
        diag1_ok = lf_swing != rh_swing  # LF与RH应相反
        diag2_ok = rf_swing != lh_swing  # RF与LH应相反
        trot_reward = 1.0 if (diag1_ok and diag2_ok) else -0.5
        
        reward = vel_reward + posture_penalty + smooth_penalty + trot_reward
        
        # 终止条件:摔倒(姿态角速度过大)
        done = abs(self.body_pitch_rate) > 0.5 or abs(self.body_roll_rate) > 0.5 or self.step_count > 500
        return self._get_obs(), reward, done, {'vel': self.forward_vel, 'trot': (diag1_ok and diag2_ok)}

# ---------- 策略网络 ----------
class PolicyNetwork:
    def __init__(self, obs_dim, act_dim, hidden=32, lr=0.01):
        # 两层网络:obs_dim → hidden → act_dim
        self.w1 = np.random.randn(obs_dim, hidden) * 0.1
        self.b1 = np.zeros(hidden)
        self.w2 = np.random.randn(hidden, act_dim) * 0.1
        self.b2 = np.zeros(act_dim)
        self.lr = lr
        self.act_dim = act_dim
        self.obs_dim = obs_dim
        self.log_std = np.zeros(act_dim)  # 固定标准差
        
    def forward(self, obs):
        h = np.tanh(obs @ self.w1 + self.b1)
        mu = h @ self.w2 + self.b2
        return mu
    
    def sample_action(self, obs):
        mu = self.forward(obs)
        std = np.exp(self.log_std)
        action = np.random.normal(mu, std)
        return np.clip(action, -1, 1), mu, std
    
    def log_prob(self, obs, action, mu, std):
        # 高斯对数概率
        return -0.5 * np.sum(((action - mu) / std)**2 + np.log(2*np.pi*std**2))
    
    def update(self, obs_history, action_history, mu_history, std_history, returns):
        """REINFORCE更新:∇θ log π(a|s) * R"""
        grad_w1 = np.zeros_like(self.w1)
        grad_b1 = np.zeros_like(self.b1)
        grad_w2 = np.zeros_like(self.w2)
        grad_b2 = np.zeros_like(self.b2)
        
        for t in range(len(obs_history)):
            obs = obs_history[t]
            action = action_history[t]
            mu = mu_history[t]
            std = std_history[t]
            G = returns[t]
            
            # 计算 log π 对网络参数的梯度(反向传播)
            h = np.tanh(obs @ self.w1 + self.b1)
            dh = 1 - h**2
            
            # d log_prob / d mu = (action - mu) / std^2
            dlogp_dmu = (action - mu) / (std**2)
            
            # 通过网络反向传播
            grad_w2 += np.outer(h, dlogp_dmu) * G
            grad_b2 += dlogp_dmu * G
            
            # 传播到第一层
            delta_hidden = (dlogp_dmu @ self.w2.T) * dh
            grad_w1 += np.outer(obs, delta_hidden) * G
            grad_b1 += delta_hidden * G
        
        # 参数更新
        self.w1 += self.lr * grad_w1 / len(obs_history)
        self.b1 += self.lr * grad_b1 / len(obs_history)
        self.w2 += self.lr * grad_w2 / len(obs_history)
        self.b2 += self.lr * grad_b2 / len(obs_history)

# ---------- 训练循环 ----------
def train_reinforce(episodes=2000, gamma=0.99):
    env = SimpleQuadrupedEnv()
    policy = PolicyNetwork(env.obs_dim, env.action_dim)
    
    for episode in range(episodes):
        obs = env.reset()
        obs_history, action_history, mu_history, std_history, reward_history = [], [], [], [], []
        done = False
        total_reward = 0
        
        while not done:
            action, mu, std = policy.sample_action(obs)
            next_obs, reward, done, info = env.step(action)
            
            obs_history.append(obs)
            action_history.append(action)
            mu_history.append(mu)
            std_history.append(std)
            reward_history.append(reward)
            total_reward += reward
            
            obs = next_obs
        
        # 计算折扣回报
        returns = np.zeros(len(reward_history))
        G = 0
        for t in reversed(range(len(reward_history))):
            G = reward_history[t] + gamma * G
            returns[t] = G
        
        # 标准化回报(降低方差)
        if len(returns) > 1 and np.std(returns) > 1e-8:
            returns = (returns - np.mean(returns)) / (np.std(returns) + 1e-8)
        
        # 策略更新
        policy.update(obs_history, action_history, mu_history, std_history, returns)
        
        if (episode+1) % 100 == 0:
            print(f"Episode {episode+1}, Total Reward: {total_reward:.2f}, "
                  f"Avg Vel: {np.mean([r for r in reward_history if 'vel' in str(r)]) if False else info['vel']:.3f}, "
                  f"Trot: {info['trot']}")
    
    return policy

# ---------- 运行入口 ----------
if __name__ == "__main__":
    print("开始REINFORCE小跑步态训练...")
    trained_policy = train_reinforce(episodes=500)
    print("训练完成")
    
    # 简单验证:运行一个回合观察步态
    env = SimpleQuadrupedEnv()
    obs = env.reset()
    done = False
    trot_count = 0
    total_steps = 0
    while not done:
        action, _, _ = trained_policy.sample_action(obs)
        obs, _, done, info = env.step(action)
        if info['trot']:
            trot_count += 1
        total_steps += 1
        if total_steps > 300:
            break
    print(f"验证回合: {total_steps}步, 小跑正确步态比例: {trot_count/total_steps:.2%}")

运行与验证

  1. 运行方式:将上述代码保存为 reinforce_trot.py,执行 python reinforce_trot.py。训练500回合约需1-2分钟(纯CPU)。
  2. 观察指标:训练过程中每100回合输出总奖励与前进速度。奖励应总体上升,速度应逐渐增大。
  3. 通过判定标准:训练结束后自动运行验证回合,若小跑正确步态比例(对角腿反相)超过70%,且前进速度大于0.3,则判定训练成功。若未达标,可增加训练回合数至2000或调整学习率。

关键说明

  • 环境为高度简化的运动学模型,不涉及真实物理接触力计算,但保留了小跑步态的核心约束(对角腿相位耦合)。
  • 策略网络仅输出动作均值,方差固定为1.0(经log_std=0),这是REINFORCE的常见简化处理。
  • 奖励中的小跑节奏项是引导策略形成对角步态的关键信号,若去掉该项,策略将退化为随机行走。
REINFORCE小跑步态训练流程 环境初始化 关节角+相位 策略采样 π(a|s,θ) 环境交互 执行动作→奖励 轨迹存储 (s, a, r) 序列 折扣回报计算 G_t = Σ γ^k r_{t+k} 策略梯度更新 θ ← θ + α∇logπ·G 虚线箭头表示下一回合循环