理论 基于模型的策略优化
修改了警告和报错
运行环境
Debian GNU/Linux 12
Python 3.9.19
torch 2.0.1
gym 0.26.2运行代码
MBPO.py
#!/usr/bin/env python
import gym
from collections import namedtuple
import itertools
from itertools import count
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions.normal import Normal
import numpy as np
import collections
import random
import matplotlib.pyplot as plt
class PolicyNet(torch.nn.Module):
def __init__(self, state_dim, hidden_dim, action_dim, action_bound):
super(PolicyNet, self).__init__()
self.fc1 = torch.nn.Linear(state_dim, hidden_dim)
self.fc_mu = torch.nn.Linear(hidden_dim, action_dim)
self.fc_std = torch.nn.Linear(hidden_dim, action_dim)
self.action_bound = action_bound
def forward(self, x):
x = F.relu(self.fc1(x))
mu = self.fc_mu(x)
std = F.softplus(self.fc_std(x))
dist = Normal(mu, std)
normal_sample = dist.rsample() # rsample()是重参数化采样函数
log_prob = dist.log_prob(normal_sample)
action = torch.tanh(normal_sample) # 计算tanh_normal分布的对数概率密度
log_prob = log_prob - torch.log(1 - torch.tanh(action).pow(2) + 1e-7)
action = action * self.action_bound
return action, log_prob
class QValueNet(torch.nn.Module):
def __init__(self, state_dim, hidden_dim, action_dim):
super(QValueNet, self).__init__()
self.fc1 = torch.nn.Linear(state_dim + action_dim, hidden_dim)
self.fc2 = torch.nn.Linear(hidden_dim, 1)
def forward(self, x, a):
cat = torch.cat([x, a], dim=1) # 拼接状态和动作
x = F.relu(self.fc1(cat))
return self.fc2(x)
device = torch.device("cuda") if torch.cuda.is_available() else torch.device(
"cpu")
class SAC:
''' 处理连续动作的SAC算法 '''
def __init__(self, state_dim, hidden_dim, action_dim, action_bound,
actor_lr, critic_lr, alpha_lr, target_entropy, tau, gamma):
self.actor = PolicyNet(state_dim, hidden_dim, action_dim,
action_bound).to(device) # 策略网络
# 第一个Q网络
self.critic_1 = QValueNet(state_dim, hidden_dim, action_dim).to(device)
# 第二个Q网络
self.critic_2 = QValueNet(state_dim, hidden_dim, action_dim).to(device)
self.target_critic_1 = QValueNet(state_dim, hidden_dim,
action_dim).to(device) # 第一个目标Q网络
self.target_critic_2 = QValueNet(state_dim, hidden_dim,
action_dim).to(device) # 第二个目标Q网络
# 令目标Q网络的初始参数和Q网络一样
self.target_critic_1.load_state_dict(self.critic_1.state_dict())
self.target_critic_2.load_state_dict(self.critic_2.state_dict())
self.actor_optimizer = torch.optim.Adam(self.actor.parameters(),
lr=actor_lr)
self.critic_1_optimizer = torch.optim.Adam(self.critic_1.parameters(),
lr=critic_lr)
self.critic_2_optimizer = torch.optim.Adam(self.critic_2.parameters(),
lr=critic_lr)
# 使用alpha的log值,可以使训练结果比较稳定
self.log_alpha = torch.tensor(np.log(0.01), dtype=torch.float)
self.log_alpha.requires_grad = True # 可以对alpha求梯度
self.log_alpha_optimizer = torch.optim.Adam([self.log_alpha],
lr=alpha_lr)
self.target_entropy = target_entropy # 目标熵的大小
self.gamma = gamma
self.tau = tau
def take_action(self, state):
state = torch.tensor(np.array([state]), dtype=torch.float).to(device)
action = self.actor(state)[0]
return [action.item()]
def calc_target(self, rewards, next_states, dones): # 计算目标Q值
next_actions, log_prob = self.actor(next_states)
entropy = -log_prob
q1_value = self.target_critic_1(next_states, next_actions)
q2_value = self.target_critic_2(next_states, next_actions)
next_value = torch.min(q1_value,
q2_value) + self.log_alpha.exp() * entropy
td_target = rewards + self.gamma * next_value * (1 - dones)
return td_target
def soft_update(self, net, target_net):
for param_target, param in zip(target_net.parameters(),
net.parameters()):
param_target.data.copy_(param_target.data * (1.0 - self.tau) +
param.data * self.tau)
def update(self, transition_dict):
states = torch.tensor(transition_dict['states'],
dtype=torch.float).to(device)
actions = torch.tensor(transition_dict['actions'],
dtype=torch.float).view(-1, 1).to(device)
rewards = torch.tensor(transition_dict['rewards'],
dtype=torch.float).view(-1, 1).to(device)
next_states = torch.tensor(transition_dict['next_states'],
dtype=torch.float).to(device)
dones = torch.tensor(transition_dict['dones'],
dtype=torch.float).view(-1, 1).to(device)
rewards = (rewards + 8.0) / 8.0 # 对倒立摆环境的奖励进行重塑
# 更新两个Q网络
td_target = self.calc_target(rewards, next_states, dones)
critic_1_loss = torch.mean(
F.mse_loss(self.critic_1(states, actions), td_target.detach()))
critic_2_loss = torch.mean(
F.mse_loss(self.critic_2(states, actions), td_target.detach()))
self.critic_1_optimizer.zero_grad()
critic_1_loss.backward()
self.critic_1_optimizer.step()
self.critic_2_optimizer.zero_grad()
critic_2_loss.backward()
self.critic_2_optimizer.step()
# 更新策略网络
new_actions, log_prob = self.actor(states)
entropy = -log_prob
q1_value = self.critic_1(states, new_actions)
q2_value = self.critic_2(states, new_actions)
actor_loss = torch.mean(-self.log_alpha.exp() * entropy -
torch.min(q1_value, q2_value))
self.actor_optimizer.zero_grad()
actor_loss.backward()
self.actor_optimizer.step()
# 更新alpha值
alpha_loss = torch.mean(
(entropy - target_entropy).detach() * self.log_alpha.exp())
self.log_alpha_optimizer.zero_grad()
alpha_loss.backward()
self.log_alpha_optimizer.step()
self.soft_update(self.critic_1, self.target_critic_1)
self.soft_update(self.critic_2, self.target_critic_2)
class Swish(nn.Module):
''' Swish激活函数 '''
def __init__(self):
super(Swish, self).__init__()
def forward(self, x):
return x * torch.sigmoid(x)
def init_weights(m):
''' 初始化模型权重 '''
def truncated_normal_init(t, mean=0.0, std=0.01):
torch.nn.init.normal_(t, mean=mean, std=std)
while True:
cond = (t < mean - 2 * std) | (t > mean + 2 * std)
if not torch.sum(cond):
break
t = torch.where(
cond,
torch.nn.init.normal_(torch.ones(t.shape, device=device),
mean=mean,
std=std), t)
return t
if type(m) == nn.Linear or isinstance(m, FCLayer):
truncated_normal_init(m.weight, std=1 / (2 * np.sqrt(m._input_dim)))
m.bias.data.fill_(0.0)
class FCLayer(nn.Module):
''' 集成之后的全连接层 '''
def __init__(self, input_dim, output_dim, ensemble_size, activation):
super(FCLayer, self).__init__()
self._input_dim, self._output_dim = input_dim, output_dim
self.weight = nn.Parameter(
torch.Tensor(ensemble_size, input_dim, output_dim).to(device))
self._activation = activation
self.bias = nn.Parameter(
torch.Tensor(ensemble_size, output_dim).to(device))
def forward(self, x):
return self._activation(
torch.add(torch.bmm(x, self.weight), self.bias[:, None, :]))
class EnsembleModel(nn.Module):
''' 环境模型集成 '''
def __init__(self,
state_dim,
action_dim,
model_alpha,
ensemble_size=5,
learning_rate=1e-3):
super(EnsembleModel, self).__init__()
# 输出包括均值和方差,因此是状态与奖励维度之和的两倍
self._output_dim = (state_dim + 1) * 2
self._model_alpha = model_alpha # 模型损失函数中加权时的权重
self._max_logvar = nn.Parameter((torch.ones(
(1, self._output_dim // 2)).float() / 2).to(device),
requires_grad=False)
self._min_logvar = nn.Parameter((-torch.ones(
(1, self._output_dim // 2)).float() * 10).to(device),
requires_grad=False)
self.layer1 = FCLayer(state_dim + action_dim, 200, ensemble_size,
Swish())
self.layer2 = FCLayer(200, 200, ensemble_size, Swish())
self.layer3 = FCLayer(200, 200, ensemble_size, Swish())
self.layer4 = FCLayer(200, 200, ensemble_size, Swish())
self.layer5 = FCLayer(200, self._output_dim, ensemble_size,
nn.Identity())
self.apply(init_weights) # 初始化环境模型中的参数
self.optimizer = torch.optim.Adam(self.parameters(), lr=learning_rate)
def forward(self, x, return_log_var=False):
ret = self.layer5(self.layer4(self.layer3(self.layer2(
self.layer1(x)))))
mean = ret[:, :, :self._output_dim // 2]
# 在PETS算法中,将方差控制在最小值和最大值之间
logvar = self._max_logvar - F.softplus(
self._max_logvar - ret[:, :, self._output_dim // 2:])
logvar = self._min_logvar + F.softplus(logvar - self._min_logvar)
return mean, logvar if return_log_var else torch.exp(logvar)
def loss(self, mean, logvar, labels, use_var_loss=True):
inverse_var = torch.exp(-logvar)
if use_var_loss:
mse_loss = torch.mean(torch.mean(torch.pow(mean - labels, 2) *
inverse_var,
dim=-1),
dim=-1)
var_loss = torch.mean(torch.mean(logvar, dim=-1), dim=-1)
total_loss = torch.sum(mse_loss) + torch.sum(var_loss)
else:
mse_loss = torch.mean(torch.pow(mean - labels, 2), dim=(1, 2))
total_loss = torch.sum(mse_loss)
return total_loss, mse_loss
def train(self, loss):
self.optimizer.zero_grad()
loss += self._model_alpha * torch.sum(
self._max_logvar) - self._model_alpha * torch.sum(self._min_logvar)
loss.backward()
self.optimizer.step()
class EnsembleDynamicsModel:
''' 环境模型集成,加入精细化的训练 '''
def __init__(self, state_dim, action_dim, model_alpha=0.01, num_network=5):
self._num_network = num_network
self._state_dim, self._action_dim = state_dim, action_dim
self.model = EnsembleModel(state_dim,
action_dim,
model_alpha,
ensemble_size=num_network)
self._epoch_since_last_update = 0
def train(self,
inputs,
labels,
batch_size=64,
holdout_ratio=0.1,
max_iter=20):
# 设置训练集与验证集
permutation = np.random.permutation(inputs.shape[0])
inputs, labels = inputs[permutation], labels[permutation]
num_holdout = int(inputs.shape[0] * holdout_ratio)
train_inputs, train_labels = inputs[num_holdout:], labels[num_holdout:]
holdout_inputs, holdout_labels = inputs[:
num_holdout], labels[:
num_holdout]
holdout_inputs = torch.from_numpy(holdout_inputs).float().to(device)
holdout_labels = torch.from_numpy(holdout_labels).float().to(device)
holdout_inputs = holdout_inputs[None, :, :].repeat(
[self._num_network, 1, 1])
holdout_labels = holdout_labels[None, :, :].repeat(
[self._num_network, 1, 1])
# 保留最好的结果
self._snapshots = {i: (None, 1e10) for i in range(self._num_network)}
for epoch in itertools.count():
# 定义每一个网络的训练数据
train_index = np.vstack([
np.random.permutation(train_inputs.shape[0])
for _ in range(self._num_network)
])
# 所有真实数据都用来训练
for batch_start_pos in range(0, train_inputs.shape[0], batch_size):
batch_index = train_index[:, batch_start_pos:batch_start_pos +
batch_size]
train_input = torch.from_numpy(
train_inputs[batch_index]).float().to(device)
train_label = torch.from_numpy(
train_labels[batch_index]).float().to(device)
mean, logvar = self.model(train_input, return_log_var=True)
loss, _ = self.model.loss(mean, logvar, train_label)
self.model.train(loss)
with torch.no_grad():
mean, logvar = self.model(holdout_inputs, return_log_var=True)
_, holdout_losses = self.model.loss(mean,
logvar,
holdout_labels,
use_var_loss=False)
holdout_losses = holdout_losses.cpu()
break_condition = self._save_best(epoch, holdout_losses)
if break_condition or epoch > max_iter: # 结束训练
break
def _save_best(self, epoch, losses, threshold=0.1):
updated = False
for i in range(len(losses)):
current = losses[i]
_, best = self._snapshots[i]
improvement = (best - current) / best
if improvement > threshold:
self._snapshots[i] = (epoch, current)
updated = True
self._epoch_since_last_update = 0 if updated else self._epoch_since_last_update + 1
return self._epoch_since_last_update > 5
def predict(self, inputs, batch_size=64):
inputs = np.tile(inputs, (self._num_network, 1, 1))
inputs = torch.tensor(inputs, dtype=torch.float).to(device)
mean, var = self.model(inputs, return_log_var=False)
return mean.detach().cpu().numpy(), var.detach().cpu().numpy()
class FakeEnv:
def __init__(self, model):
self.model = model
def step(self, obs, act):
inputs = np.concatenate((obs, act), axis=-1)
ensemble_model_means, ensemble_model_vars = self.model.predict(inputs)
ensemble_model_means[:, :, 1:] += obs
ensemble_model_stds = np.sqrt(ensemble_model_vars)
ensemble_samples = ensemble_model_means + np.random.normal(
size=ensemble_model_means.shape) * ensemble_model_stds
num_models, batch_size, _ = ensemble_model_means.shape
models_to_use = np.random.choice(
[i for i in range(self.model._num_network)], size=batch_size)
batch_inds = np.arange(0, batch_size)
samples = ensemble_samples[models_to_use, batch_inds]
rewards, next_obs = samples[:, :1][0][0], samples[:, 1:][0]
return rewards, next_obs
class MBPO:
def __init__(self, env, agent, fake_env, env_pool, model_pool,
rollout_length, rollout_batch_size, real_ratio, num_episode):
self.env = env
self.agent = agent
self.fake_env = fake_env
self.env_pool = env_pool
self.model_pool = model_pool
self.rollout_length = rollout_length
self.rollout_batch_size = rollout_batch_size
self.real_ratio = real_ratio
self.num_episode = num_episode
def rollout_model(self):
observations, _, _, _, _ = self.env_pool.sample(
self.rollout_batch_size)
for obs in observations:
for i in range(self.rollout_length):
action = self.agent.take_action(obs)
reward, next_obs = self.fake_env.step(obs, action)
self.model_pool.add(obs, action, reward, next_obs, False)
obs = next_obs
def update_agent(self, policy_train_batch_size=64):
env_batch_size = int(policy_train_batch_size * self.real_ratio)
model_batch_size = policy_train_batch_size - env_batch_size
for epoch in range(10):
env_obs, env_action, env_reward, env_next_obs, env_done = self.env_pool.sample(
env_batch_size)
if self.model_pool.size() > 0:
model_obs, model_action, model_reward, model_next_obs, model_done = self.model_pool.sample(
model_batch_size)
obs = np.concatenate((env_obs, model_obs), axis=0)
action = np.concatenate((env_action, model_action), axis=0)
next_obs = np.concatenate((env_next_obs, model_next_obs),
axis=0)
reward = np.concatenate((env_reward, model_reward), axis=0)
done = np.concatenate((env_done, model_done), axis=0)
else:
obs, action, next_obs, reward, done = env_obs, env_action, env_next_obs, env_reward, env_done
transition_dict = {
'states': obs,
'actions': action,
'next_states': next_obs,
'rewards': reward,
'dones': done
}
self.agent.update(transition_dict)
def train_model(self):
obs, action, reward, next_obs, done = self.env_pool.return_all_samples(
)
inputs = np.concatenate((obs, action), axis=-1)
reward = np.array(reward)
labels = np.concatenate(
(np.reshape(reward, (reward.shape[0], -1)), next_obs - obs),
axis=-1)
self.fake_env.model.train(inputs, labels)
def explore(self):
obs, done, episode_return = self.env.reset()[0], False, 0
num = 0
while not done and num < 10000:
action = self.agent.take_action(obs)
next_obs, reward, done, _, __ = self.env.step(action)
self.env_pool.add(obs, action, reward, next_obs, done)
obs = next_obs
episode_return += reward
num = num + 1
return episode_return
def train(self):
return_list = []
explore_return = self.explore() # 随机探索采取数据
print('episode: 1, return: %d' % explore_return)
return_list.append(explore_return)
for i_episode in range(self.num_episode - 1):
obs, done, episode_return = self.env.reset()[0], False, 0
step = 0
while not done:
if step % 50 == 0:
print(step)
self.train_model()
self.rollout_model()
if step >= 100:
break
action = self.agent.take_action(obs)
next_obs, reward, done, _, __ = self.env.step(action)
self.env_pool.add(obs, action, reward, next_obs, done)
obs = next_obs
episode_return += reward
self.update_agent()
step += 1
return_list.append(episode_return)
print('episode: %d, return: %d' % (i_episode + 2, episode_return))
return return_list
class ReplayBuffer:
def __init__(self, capacity):
self.buffer = collections.deque(maxlen=capacity)
def add(self, state, action, reward, next_state, done):
self.buffer.append((state, action, reward, next_state, done))
def size(self):
return len(self.buffer)
def sample(self, batch_size):
if batch_size > len(self.buffer):
return self.return_all_samples()
else:
transitions = random.sample(self.buffer, batch_size)
state, action, reward, next_state, done = zip(*transitions)
return np.array(state), action, reward, np.array(next_state), done
def return_all_samples(self):
all_transitions = list(self.buffer)
state, action, reward, next_state, done = zip(*all_transitions)
return np.array(state), action, reward, np.array(next_state), done
real_ratio = 0.5
env_name = 'Pendulum-v1'
env = gym.make(env_name)
num_episodes = 20
actor_lr = 5e-4
critic_lr = 5e-3
alpha_lr = 1e-3
hidden_dim = 128
gamma = 0.98
tau = 0.005 # 软更新参数
buffer_size = 10000
target_entropy = -1
model_alpha = 0.01 # 模型损失函数中的加权权重
state_dim = env.observation_space.shape[0]
action_dim = env.action_space.shape[0]
action_bound = env.action_space.high[0] # 动作最大值
rollout_batch_size = 1000
rollout_length = 1 # 推演长度k,推荐更多尝试
model_pool_size = rollout_batch_size * rollout_length
agent = SAC(state_dim, hidden_dim, action_dim, action_bound, actor_lr,
critic_lr, alpha_lr, target_entropy, tau, gamma)
model = EnsembleDynamicsModel(state_dim, action_dim, model_alpha)
fake_env = FakeEnv(model)
env_pool = ReplayBuffer(buffer_size)
model_pool = ReplayBuffer(model_pool_size)
mbpo = MBPO(env, agent, fake_env, env_pool, model_pool, rollout_length,
rollout_batch_size, real_ratio, num_episodes)
return_list = mbpo.train()
episodes_list = list(range(len(return_list)))
plt.plot(episodes_list, return_list)
plt.xlabel('Episodes')
plt.ylabel('Returns')
plt.title('MBPO on {}'.format(env_name))
plt.show()
478
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
