wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-07-21 12:50:45 +08:00
Code Issues Packages Projects Releases Wiki Activity

[rllib] Refactor pytorch custom model support (#3634)

Browse Source
This commit is contained in:
Eric Liang
2019-01-03 13:48:33 +08:00
committed by GitHub
parent b6bcd18d65
commit 47d36d7bd6
19 changed files with 402 additions and 240 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/ray/rllib/agents/a3c/a3c_torch_policy_graph.py
+6 -9
View File
@@ -7,7 +7,6 @@ import torch.nn.functional as F
from torch import nn
import ray
from ray.rllib.models.pytorch.misc import var_to_np
from ray.rllib.models.catalog import ModelCatalog
from ray.rllib.evaluation.postprocessing import compute_advantages
from ray.rllib.evaluation.policy_graph import PolicyGraph
@@ -23,7 +22,7 @@ class A3CLoss(nn.Module):
self.entropy_coeff = entropy_coeff
def forward(self, observations, actions, advantages, value_targets):
logits, values = self.policy_model(observations)
logits, _, values, _ = self.policy_model({"obs": observations}, [])
log_probs = F.log_softmax(logits, dim=1)
probs = F.softmax(logits, dim=1)
action_log_probs = log_probs.gather(1, actions.view(-1, 1))
@@ -46,8 +45,8 @@ class A3CTorchPolicyGraph(TorchPolicyGraph):
self.config = config
_, self.logit_dim = ModelCatalog.get_action_dist(
action_space, self.config["model"])
self.model = ModelCatalog.get_torch_model(
obs_space.shape, self.logit_dim, self.config["model"])
self.model = ModelCatalog.get_torch_model(obs_space, self.logit_dim,
self.config["model"])
loss = A3CLoss(self.model, self.config["vf_loss_coeff"],
self.config["entropy_coeff"])
TorchPolicyGraph.__init__(
@@ -60,7 +59,7 @@ class A3CTorchPolicyGraph(TorchPolicyGraph):
@override(TorchPolicyGraph)
def extra_action_out(self, model_out):
return {"vf_preds": var_to_np(model_out[1])}
return {"vf_preds": model_out[2].numpy()}
@override(TorchPolicyGraph)
def optimizer(self):
@@ -82,7 +81,5 @@ class A3CTorchPolicyGraph(TorchPolicyGraph):
def _value(self, obs):
with self.lock:
obs = torch.from_numpy(obs).float().unsqueeze(0)
res = self.model.hidden_layers(obs)
res = self.model.value_branch(res)
res = res.squeeze()
return var_to_np(res)
_, _, vf, _ = self.model({"obs": obs}, [])
return vf.numpy().squeeze()
python/ray/rllib/agents/qmix/model.py
+26 -15
View File
@@ -5,24 +5,35 @@ from __future__ import print_function
from torch import nn
import torch.nn.functional as F
from ray.rllib.models.preprocessors import get_preprocessor
from ray.rllib.models.pytorch.model import TorchModel
from ray.rllib.utils.annotations import override
# TODO(ekl) we should have common models for pytorch like we do for TF
class RNNModel(nn.Module):
def __init__(self, obs_size, rnn_hidden_dim, n_actions):
nn.Module.__init__(self)
self.rnn_hidden_dim = rnn_hidden_dim
self.n_actions = n_actions
self.fc1 = nn.Linear(obs_size, rnn_hidden_dim)
self.rnn = nn.GRUCell(rnn_hidden_dim, rnn_hidden_dim)
self.fc2 = nn.Linear(rnn_hidden_dim, n_actions)
def init_hidden(self):
class RNNModel(TorchModel):
"""The default RNN model for QMIX."""
def __init__(self, obs_space, num_outputs, options):
TorchModel.__init__(self, obs_space, num_outputs, options)
self.obs_size = _get_size(obs_space)
self.rnn_hidden_dim = options["lstm_cell_size"]
self.fc1 = nn.Linear(self.obs_size, self.rnn_hidden_dim)
self.rnn = nn.GRUCell(self.rnn_hidden_dim, self.rnn_hidden_dim)
self.fc2 = nn.Linear(self.rnn_hidden_dim, num_outputs)
@override(TorchModel)
def state_init(self):
# make hidden states on same device as model
return self.fc1.weight.new(1, self.rnn_hidden_dim).zero_()
return [self.fc1.weight.new(1, self.rnn_hidden_dim).zero_().squeeze(0)]
def forward(self, inputs, hidden_state):
x = F.relu(self.fc1(inputs.float()))
h_in = hidden_state.reshape(-1, self.rnn_hidden_dim)
@override(TorchModel)
def _forward(self, input_dict, hidden_state):
x = F.relu(self.fc1(input_dict["obs"]))
h_in = hidden_state[0].reshape(-1, self.rnn_hidden_dim)
h = self.rnn(x, h_in)
q = self.fc2(h)
return q, h
return q, h, None, [h]
def _get_size(obs_space):
return get_preprocessor(obs_space)(obs_space).size
python/ray/rllib/agents/qmix/qmix_policy_graph.py
+40 -29
View File
@@ -12,13 +12,12 @@ from torch.distributions import Categorical
import ray
from ray.rllib.agents.qmix.mixers import VDNMixer, QMixer
from ray.rllib.agents.qmix.model import RNNModel
from ray.rllib.agents.qmix.model import RNNModel, _get_size
from ray.rllib.evaluation.policy_graph import PolicyGraph
from ray.rllib.models.action_dist import TupleActions
from ray.rllib.models.pytorch.misc import var_to_np
from ray.rllib.models.catalog import ModelCatalog
from ray.rllib.models.lstm import chop_into_sequences
from ray.rllib.models.model import _unpack_obs
from ray.rllib.models.preprocessors import get_preprocessor
from ray.rllib.env.constants import GROUP_REWARDS
from ray.rllib.utils.annotations import override
@@ -61,7 +60,7 @@ class QMixLoss(nn.Module):
# Calculate estimated Q-Values
mac_out = []
h = self.model.init_hidden().expand([B, self.n_agents, -1])
h = [s.expand([B, self.n_agents, -1]) for s in self.model.state_init()]
for t in range(T):
q, h = _mac(self.model, obs[:, t], h)
mac_out.append(q)
@@ -73,8 +72,10 @@ class QMixLoss(nn.Module):
# Calculate the Q-Values necessary for the target
target_mac_out = []
target_h = self.target_model.init_hidden().expand(
[B, self.n_agents, -1])
target_h = [
s.expand([B, self.n_agents, -1])
for s in self.target_model.state_init()
]
for t in range(T):
target_q, target_h = _mac(self.target_model, obs[:, t], target_h)
target_mac_out.append(target_q)
@@ -154,13 +155,22 @@ class QMixPolicyGraph(PolicyGraph):
(self.n_actions, ), mask_shape))
self.has_action_mask = True
self.obs_size = _get_size(agent_obs_space.spaces["obs"])
# The real agent obs space is nested inside the dict
agent_obs_space = agent_obs_space.spaces["obs"]
else:
self.has_action_mask = False
self.obs_size = _get_size(agent_obs_space)
self.model = RNNModel(self.obs_size, self.h_size, self.n_actions)
self.target_model = RNNModel(self.obs_size, self.h_size,
self.n_actions)
self.model = ModelCatalog.get_torch_model(
agent_obs_space,
self.n_actions,
config["model"],
default_model_cls=RNNModel)
self.target_model = ModelCatalog.get_torch_model(
agent_obs_space,
self.n_actions,
config["model"],
default_model_cls=RNNModel)
# Setup the mixer network.
# The global state is just the stacked agent observations for now.
@@ -203,13 +213,12 @@ class QMixPolicyGraph(PolicyGraph):
episodes=None,
**kwargs):
obs_batch, action_mask = self._unpack_observation(obs_batch)
assert len(state_batches) == self.n_agents, state_batches
state_batches = np.stack(state_batches, axis=1)
# Compute actions
with th.no_grad():
q_values, hiddens = _mac(self.model, th.from_numpy(obs_batch),
th.from_numpy(state_batches))
q_values, hiddens = _mac(
self.model, th.from_numpy(obs_batch),
[th.from_numpy(np.array(s)) for s in state_batches])
avail = th.from_numpy(action_mask).float()
masked_q_values = q_values.clone()
masked_q_values[avail == 0.0] = -float("inf")
@@ -219,11 +228,10 @@ class QMixPolicyGraph(PolicyGraph):
random_actions = Categorical(avail).sample().long()
actions = (pick_random * random_actions +
(1 - pick_random) * masked_q_values.max(dim=2)[1])
actions = var_to_np(actions)
hiddens = var_to_np(hiddens)
actions = actions.numpy()
hiddens = [s.numpy() for s in hiddens]
return (TupleActions(list(actions.transpose([1, 0]))),
hiddens.transpose([1, 0, 2]), {})
return TupleActions(list(actions.transpose([1, 0]))), hiddens, {}
@override(PolicyGraph)
def compute_apply(self, samples):
@@ -239,7 +247,7 @@ class QMixPolicyGraph(PolicyGraph):
samples["dones"], obs_batch
],
[samples["state_in_{}".format(k)]
for k in range(self.n_agents)],
for k in range(len(self.get_initial_state()))],
max_seq_len=self.config["model"]["max_seq_len"],
dynamic_max=True,
_extra_padding=1)
@@ -292,8 +300,8 @@ class QMixPolicyGraph(PolicyGraph):
@override(PolicyGraph)
def get_initial_state(self):
return [
self.model.init_hidden().numpy().squeeze()
for _ in range(self.n_agents)
s.expand([self.n_agents, -1]).numpy()
for s in self.model.state_init()
]
@override(PolicyGraph)
@@ -342,6 +350,12 @@ class QMixPolicyGraph(PolicyGraph):
return group_rewards
def _unpack_observation(self, obs_batch):
"""Unpacks the action mask / tuple obs from agent grouping.
Returns:
obs (Tensor): flattened obs tensor of shape [B, n_agents, obs_size]
mask (Tensor): action mask, if any
"""
unpacked = _unpack_obs(
np.array(obs_batch),
self.observation_space.original_space,
@@ -388,17 +402,13 @@ def _validate(obs_space, action_space):
"must be homogeneous, got {}".format(action_space.spaces))
def _get_size(obs_space):
return get_preprocessor(obs_space)(obs_space).size
def _mac(model, obs, h):
"""Forward pass of the multi-agent controller.
Arguments:
model: Model that produces q-values for a 1d agent batch
model: TorchModel class
obs: Tensor of shape [B, n_agents, obs_size]
h: Tensor of shape [B, n_agents, h_size]
h: List of tensors of shape [B, n_agents, h_size]
Returns:
q_vals: Tensor of shape [B, n_agents, n_actions]
@@ -406,6 +416,7 @@ def _mac(model, obs, h):
"""
B, n_agents = obs.size(0), obs.size(1)
obs_flat = obs.reshape([B * n_agents, -1])
h_flat = h.reshape([B * n_agents, -1])
q_flat, h_flat = model.forward(obs_flat, h_flat)
return q_flat.reshape([B, n_agents, -1]), h_flat.reshape([B, n_agents, -1])
h_flat = [s.reshape([B * n_agents, -1]) for s in h]
q_flat, _, _, h_flat = model.forward({"obs": obs_flat}, h_flat)
return q_flat.reshape(
[B, n_agents, -1]), [s.reshape([B, n_agents, -1]) for s in h_flat]