mirror of
https://github.com/vale981/ray
synced 2025-03-06 10:31:39 -05:00
430 lines
16 KiB
Python
430 lines
16 KiB
Python
import logging
|
|
|
|
import ray
|
|
from ray.rllib.evaluation.postprocessing import Postprocessing
|
|
from ray.rllib.policy.sample_batch import SampleBatch
|
|
from ray.rllib.policy.torch_policy_template import build_torch_policy
|
|
from ray.rllib.agents.ppo.ppo_tf_policy import postprocess_ppo_gae, \
|
|
setup_config
|
|
from ray.rllib.agents.ppo.ppo_torch_policy import vf_preds_fetches, \
|
|
ValueNetworkMixin
|
|
from ray.rllib.agents.a3c.a3c_torch_policy import apply_grad_clipping
|
|
from ray.rllib.utils.framework import get_activation_fn
|
|
from ray.rllib.utils.framework import try_import_torch
|
|
|
|
torch, nn = try_import_torch()
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def PPOLoss(dist_class,
|
|
actions,
|
|
curr_logits,
|
|
behaviour_logits,
|
|
advantages,
|
|
value_fn,
|
|
value_targets,
|
|
vf_preds,
|
|
cur_kl_coeff,
|
|
entropy_coeff,
|
|
clip_param,
|
|
vf_clip_param,
|
|
vf_loss_coeff,
|
|
clip_loss=False):
|
|
def surrogate_loss(actions, curr_dist, prev_dist, advantages, clip_param,
|
|
clip_loss):
|
|
pi_new_logp = curr_dist.logp(actions)
|
|
pi_old_logp = prev_dist.logp(actions)
|
|
|
|
logp_ratio = torch.exp(pi_new_logp - pi_old_logp)
|
|
if clip_loss:
|
|
return torch.min(
|
|
advantages * logp_ratio,
|
|
advantages * torch.clamp(logp_ratio, 1 - clip_param,
|
|
1 + clip_param))
|
|
return advantages * logp_ratio
|
|
|
|
def kl_loss(curr_dist, prev_dist):
|
|
return prev_dist.kl(curr_dist)
|
|
|
|
def entropy_loss(dist):
|
|
return dist.entropy()
|
|
|
|
def vf_loss(value_fn, value_targets, vf_preds, vf_clip_param=0.1):
|
|
# GAE Value Function Loss
|
|
vf_loss1 = torch.pow(value_fn - value_targets, 2.0)
|
|
vf_clipped = vf_preds + torch.clamp(value_fn - vf_preds,
|
|
-vf_clip_param, vf_clip_param)
|
|
vf_loss2 = torch.pow(vf_clipped - value_targets, 2.0)
|
|
vf_loss = torch.max(vf_loss1, vf_loss2)
|
|
return vf_loss
|
|
|
|
pi_new_dist = dist_class(curr_logits, None)
|
|
pi_old_dist = dist_class(behaviour_logits, None)
|
|
|
|
surr_loss = torch.mean(
|
|
surrogate_loss(actions, pi_new_dist, pi_old_dist, advantages,
|
|
clip_param, clip_loss))
|
|
kl_loss = torch.mean(kl_loss(pi_new_dist, pi_old_dist))
|
|
vf_loss = torch.mean(
|
|
vf_loss(value_fn, value_targets, vf_preds, vf_clip_param))
|
|
entropy_loss = torch.mean(entropy_loss(pi_new_dist))
|
|
|
|
total_loss = -surr_loss + cur_kl_coeff * kl_loss
|
|
total_loss += vf_loss_coeff * vf_loss
|
|
total_loss -= entropy_coeff * entropy_loss
|
|
return total_loss, surr_loss, kl_loss, vf_loss, entropy_loss
|
|
|
|
|
|
# This is the computation graph for workers (inner adaptation steps)
|
|
class WorkerLoss(object):
|
|
def __init__(self,
|
|
model,
|
|
dist_class,
|
|
actions,
|
|
curr_logits,
|
|
behaviour_logits,
|
|
advantages,
|
|
value_fn,
|
|
value_targets,
|
|
vf_preds,
|
|
cur_kl_coeff,
|
|
entropy_coeff,
|
|
clip_param,
|
|
vf_clip_param,
|
|
vf_loss_coeff,
|
|
clip_loss=False):
|
|
self.loss, surr_loss, kl_loss, vf_loss, ent_loss = PPOLoss(
|
|
dist_class=dist_class,
|
|
actions=actions,
|
|
curr_logits=curr_logits,
|
|
behaviour_logits=behaviour_logits,
|
|
advantages=advantages,
|
|
value_fn=value_fn,
|
|
value_targets=value_targets,
|
|
vf_preds=vf_preds,
|
|
cur_kl_coeff=cur_kl_coeff,
|
|
entropy_coeff=entropy_coeff,
|
|
clip_param=clip_param,
|
|
vf_clip_param=vf_clip_param,
|
|
vf_loss_coeff=vf_loss_coeff,
|
|
clip_loss=clip_loss)
|
|
|
|
|
|
# This is the Meta-Update computation graph for main (meta-update step)
|
|
class MAMLLoss(object):
|
|
def __init__(self,
|
|
model,
|
|
config,
|
|
dist_class,
|
|
value_targets,
|
|
advantages,
|
|
actions,
|
|
behaviour_logits,
|
|
vf_preds,
|
|
cur_kl_coeff,
|
|
policy_vars,
|
|
obs,
|
|
num_tasks,
|
|
split,
|
|
inner_adaptation_steps=1,
|
|
entropy_coeff=0,
|
|
clip_param=0.3,
|
|
vf_clip_param=0.1,
|
|
vf_loss_coeff=1.0,
|
|
use_gae=True):
|
|
|
|
self.config = config
|
|
self.num_tasks = num_tasks
|
|
self.inner_adaptation_steps = inner_adaptation_steps
|
|
self.clip_param = clip_param
|
|
self.dist_class = dist_class
|
|
self.cur_kl_coeff = cur_kl_coeff
|
|
|
|
# Split episode tensors into [inner_adaptation_steps+1, num_tasks, -1]
|
|
self.obs = self.split_placeholders(obs, split)
|
|
self.actions = self.split_placeholders(actions, split)
|
|
self.behaviour_logits = self.split_placeholders(
|
|
behaviour_logits, split)
|
|
self.advantages = self.split_placeholders(advantages, split)
|
|
self.value_targets = self.split_placeholders(value_targets, split)
|
|
self.vf_preds = self.split_placeholders(vf_preds, split)
|
|
|
|
# Construct name to tensor dictionary for easier indexing
|
|
self.policy_vars = {}
|
|
for name, w in policy_vars:
|
|
self.policy_vars[name] = w
|
|
|
|
# Calculate pi_new for PPO
|
|
pi_new_logits, current_policy_vars, value_fns = [], [], []
|
|
for i in range(self.num_tasks):
|
|
pi_new, value_fn = self.feed_forward(
|
|
self.obs[0][i],
|
|
self.policy_vars,
|
|
policy_config=config["model"])
|
|
pi_new_logits.append(pi_new)
|
|
value_fns.append(value_fn)
|
|
current_policy_vars.append(self.policy_vars)
|
|
|
|
inner_kls = []
|
|
inner_ppo_loss = []
|
|
|
|
# Recompute weights for inner-adaptation (same weights as workers)
|
|
for step in range(self.inner_adaptation_steps):
|
|
kls = []
|
|
for i in range(self.num_tasks):
|
|
# PPO Loss Function (only Surrogate)
|
|
ppo_loss, _, kl_loss, _, _ = PPOLoss(
|
|
dist_class=dist_class,
|
|
actions=self.actions[step][i],
|
|
curr_logits=pi_new_logits[i],
|
|
behaviour_logits=self.behaviour_logits[step][i],
|
|
advantages=self.advantages[step][i],
|
|
value_fn=value_fns[i],
|
|
value_targets=self.value_targets[step][i],
|
|
vf_preds=self.vf_preds[step][i],
|
|
cur_kl_coeff=0.0,
|
|
entropy_coeff=entropy_coeff,
|
|
clip_param=clip_param,
|
|
vf_clip_param=vf_clip_param,
|
|
vf_loss_coeff=vf_loss_coeff,
|
|
clip_loss=False)
|
|
|
|
adapted_policy_vars = self.compute_updated_variables(
|
|
ppo_loss, current_policy_vars[i], model)
|
|
pi_new_logits[i], value_fns[i] = self.feed_forward(
|
|
self.obs[step + 1][i],
|
|
adapted_policy_vars,
|
|
policy_config=config["model"])
|
|
current_policy_vars[i] = adapted_policy_vars
|
|
kls.append(kl_loss)
|
|
inner_ppo_loss.append(ppo_loss)
|
|
inner_kls.extend(kls)
|
|
|
|
self.mean_inner_kl = inner_kls
|
|
|
|
ppo_obj = []
|
|
for i in range(self.num_tasks):
|
|
ppo_loss, surr_loss, kl_loss, val_loss, entropy_loss = PPOLoss(
|
|
dist_class=dist_class,
|
|
actions=self.actions[self.inner_adaptation_steps][i],
|
|
curr_logits=pi_new_logits[i],
|
|
behaviour_logits=self.behaviour_logits[
|
|
self.inner_adaptation_steps][i],
|
|
advantages=self.advantages[self.inner_adaptation_steps][i],
|
|
value_fn=value_fns[i],
|
|
value_targets=self.value_targets[self.inner_adaptation_steps][
|
|
i],
|
|
vf_preds=self.vf_preds[self.inner_adaptation_steps][i],
|
|
cur_kl_coeff=0.0,
|
|
entropy_coeff=entropy_coeff,
|
|
clip_param=clip_param,
|
|
vf_clip_param=vf_clip_param,
|
|
vf_loss_coeff=vf_loss_coeff,
|
|
clip_loss=True)
|
|
ppo_obj.append(ppo_loss)
|
|
self.mean_policy_loss = surr_loss
|
|
self.mean_kl = kl_loss
|
|
self.mean_vf_loss = val_loss
|
|
self.mean_entropy = entropy_loss
|
|
|
|
self.inner_kl_loss = torch.mean(
|
|
torch.stack([
|
|
a * b for a, b in zip(self.cur_kl_coeff, self.mean_inner_kl)
|
|
]))
|
|
self.loss = torch.mean(torch.stack(ppo_obj)) + self.inner_kl_loss
|
|
|
|
def feed_forward(self, obs, policy_vars, policy_config):
|
|
# Hacky for now, reconstruct FC network with adapted weights
|
|
# @mluo: TODO for any network
|
|
def fc_network(inp, network_vars, hidden_nonlinearity,
|
|
output_nonlinearity, policy_config, hiddens_name,
|
|
logits_name):
|
|
x = inp
|
|
|
|
hidden_w = []
|
|
logits_w = []
|
|
|
|
for name, w in network_vars.items():
|
|
if hiddens_name in name:
|
|
hidden_w.append(w)
|
|
elif logits_name in name:
|
|
logits_w.append(w)
|
|
else:
|
|
raise NameError
|
|
|
|
assert len(hidden_w) % 2 == 0 and len(logits_w) == 2
|
|
|
|
while len(hidden_w) != 0:
|
|
x = nn.functional.linear(x, hidden_w.pop(0), hidden_w.pop(0))
|
|
x = hidden_nonlinearity()(x)
|
|
|
|
x = nn.functional.linear(x, logits_w.pop(0), logits_w.pop(0))
|
|
x = output_nonlinearity()(x)
|
|
|
|
return x
|
|
|
|
policyn_vars = {}
|
|
valuen_vars = {}
|
|
log_std = None
|
|
for name, param in policy_vars.items():
|
|
if "value" in name:
|
|
valuen_vars[name] = param
|
|
elif "log_std" in name:
|
|
log_std = param
|
|
else:
|
|
policyn_vars[name] = param
|
|
|
|
output_nonlinearity = nn.Identity
|
|
hidden_nonlinearity = get_activation_fn(
|
|
policy_config["fcnet_activation"], framework="torch")
|
|
|
|
pi_new_logits = fc_network(obs, policyn_vars, hidden_nonlinearity,
|
|
output_nonlinearity, policy_config,
|
|
"hidden_layers", "logits")
|
|
if log_std is not None:
|
|
pi_new_logits = torch.cat(
|
|
[
|
|
pi_new_logits,
|
|
log_std.unsqueeze(0).repeat([len(pi_new_logits), 1])
|
|
],
|
|
axis=1)
|
|
|
|
value_fn = fc_network(obs, valuen_vars, hidden_nonlinearity,
|
|
output_nonlinearity, policy_config,
|
|
"value_branch_separate", "value_branch")
|
|
|
|
return pi_new_logits, torch.squeeze(value_fn)
|
|
|
|
def compute_updated_variables(self, loss, network_vars, model):
|
|
grad = torch.autograd.grad(
|
|
loss,
|
|
inputs=model.parameters(),
|
|
create_graph=True,
|
|
allow_unused=True)
|
|
adapted_vars = {}
|
|
for i, tup in enumerate(network_vars.items()):
|
|
name, var = tup
|
|
if grad[i] is None:
|
|
adapted_vars[name] = var
|
|
else:
|
|
adapted_vars[name] = var - self.config["inner_lr"] * grad[i]
|
|
return adapted_vars
|
|
|
|
def split_placeholders(self, placeholder, split):
|
|
inner_placeholder_list = torch.split(
|
|
placeholder, torch.sum(split, dim=1).tolist(), dim=0)
|
|
placeholder_list = []
|
|
for index, split_placeholder in enumerate(inner_placeholder_list):
|
|
placeholder_list.append(
|
|
torch.split(split_placeholder, split[index].tolist(), dim=0))
|
|
return placeholder_list
|
|
|
|
|
|
def maml_loss(policy, model, dist_class, train_batch):
|
|
logits, state = model.from_batch(train_batch)
|
|
policy.cur_lr = policy.config["lr"]
|
|
|
|
if policy.config["worker_index"]:
|
|
policy.loss_obj = WorkerLoss(
|
|
model=model,
|
|
dist_class=dist_class,
|
|
actions=train_batch[SampleBatch.ACTIONS],
|
|
curr_logits=logits,
|
|
behaviour_logits=train_batch[SampleBatch.ACTION_DIST_INPUTS],
|
|
advantages=train_batch[Postprocessing.ADVANTAGES],
|
|
value_fn=model.value_function(),
|
|
value_targets=train_batch[Postprocessing.VALUE_TARGETS],
|
|
vf_preds=train_batch[SampleBatch.VF_PREDS],
|
|
cur_kl_coeff=0.0,
|
|
entropy_coeff=policy.config["entropy_coeff"],
|
|
clip_param=policy.config["clip_param"],
|
|
vf_clip_param=policy.config["vf_clip_param"],
|
|
vf_loss_coeff=policy.config["vf_loss_coeff"],
|
|
clip_loss=False)
|
|
else:
|
|
policy.var_list = model.named_parameters()
|
|
|
|
policy.loss_obj = MAMLLoss(
|
|
model=model,
|
|
dist_class=dist_class,
|
|
value_targets=train_batch[Postprocessing.VALUE_TARGETS],
|
|
advantages=train_batch[Postprocessing.ADVANTAGES],
|
|
actions=train_batch[SampleBatch.ACTIONS],
|
|
behaviour_logits=train_batch[SampleBatch.ACTION_DIST_INPUTS],
|
|
vf_preds=train_batch[SampleBatch.VF_PREDS],
|
|
cur_kl_coeff=policy.kl_coeff_val,
|
|
policy_vars=policy.var_list,
|
|
obs=train_batch[SampleBatch.CUR_OBS],
|
|
num_tasks=policy.config["num_workers"],
|
|
split=train_batch["split"],
|
|
config=policy.config,
|
|
inner_adaptation_steps=policy.config["inner_adaptation_steps"],
|
|
entropy_coeff=policy.config["entropy_coeff"],
|
|
clip_param=policy.config["clip_param"],
|
|
vf_clip_param=policy.config["vf_clip_param"],
|
|
vf_loss_coeff=policy.config["vf_loss_coeff"],
|
|
use_gae=policy.config["use_gae"])
|
|
|
|
return policy.loss_obj.loss
|
|
|
|
|
|
def maml_stats(policy, train_batch):
|
|
if policy.config["worker_index"]:
|
|
return {"worker_loss": policy.loss_obj.loss}
|
|
else:
|
|
return {
|
|
"cur_kl_coeff": policy.kl_coeff_val,
|
|
"cur_lr": policy.cur_lr,
|
|
"total_loss": policy.loss_obj.loss,
|
|
"policy_loss": policy.loss_obj.mean_policy_loss,
|
|
"vf_loss": policy.loss_obj.mean_vf_loss,
|
|
"kl": policy.loss_obj.mean_kl,
|
|
"inner_kl": policy.loss_obj.mean_inner_kl,
|
|
"entropy": policy.loss_obj.mean_entropy,
|
|
}
|
|
|
|
|
|
class KLCoeffMixin:
|
|
def __init__(self, config):
|
|
self.kl_coeff_val = [
|
|
config["kl_coeff"]
|
|
] * config["inner_adaptation_steps"] * config["num_workers"]
|
|
self.kl_target = self.config["kl_target"]
|
|
|
|
def update_kls(self, sampled_kls):
|
|
for i, kl in enumerate(sampled_kls):
|
|
if kl < self.kl_target / 1.5:
|
|
self.kl_coeff_val[i] *= 0.5
|
|
elif kl > 1.5 * self.kl_target:
|
|
self.kl_coeff_val[i] *= 2.0
|
|
return self.kl_coeff_val
|
|
|
|
|
|
def maml_optimizer_fn(policy, config):
|
|
"""
|
|
Workers use simple SGD for inner adaptation
|
|
Meta-Policy uses Adam optimizer for meta-update
|
|
"""
|
|
if not config["worker_index"]:
|
|
return torch.optim.Adam(policy.model.parameters(), lr=config["lr"])
|
|
return torch.optim.SGD(policy.model.parameters(), lr=config["inner_lr"])
|
|
|
|
|
|
def setup_mixins(policy, obs_space, action_space, config):
|
|
ValueNetworkMixin.__init__(policy, obs_space, action_space, config)
|
|
KLCoeffMixin.__init__(policy, config)
|
|
|
|
|
|
MAMLTorchPolicy = build_torch_policy(
|
|
name="MAMLTorchPolicy",
|
|
get_default_config=lambda: ray.rllib.agents.maml.maml.DEFAULT_CONFIG,
|
|
loss_fn=maml_loss,
|
|
stats_fn=maml_stats,
|
|
optimizer_fn=maml_optimizer_fn,
|
|
extra_action_out_fn=vf_preds_fetches,
|
|
postprocess_fn=postprocess_ppo_gae,
|
|
extra_grad_process_fn=apply_grad_clipping,
|
|
before_init=setup_config,
|
|
after_init=setup_mixins,
|
|
mixins=[KLCoeffMixin])
|