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%}")
运行与验证
- 运行方式:将上述代码保存为
reinforce_trot.py,执行python reinforce_trot.py。训练500回合约需1-2分钟(纯CPU)。 - 观察指标:训练过程中每100回合输出总奖励与前进速度。奖励应总体上升,速度应逐渐增大。
- 通过判定标准:训练结束后自动运行验证回合,若小跑正确步态比例(对角腿反相)超过70%,且前进速度大于0.3,则判定训练成功。若未达标,可增加训练回合数至2000或调整学习率。
关键说明
- 环境为高度简化的运动学模型,不涉及真实物理接触力计算,但保留了小跑步态的核心约束(对角腿相位耦合)。
- 策略网络仅输出动作均值,方差固定为1.0(经
log_std=0),这是REINFORCE的常见简化处理。 - 奖励中的小跑节奏项是引导策略形成对角步态的关键信号,若去掉该项,策略将退化为随机行走。