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

[rllib] Fix torch TD error, IMPALA LR updates (#9477)

Browse Source 
* update

* add test

* lint

* fix super call

* speed es test up
This commit is contained in:
Eric Liang
2020-07-23 12:50:25 -07:00
committed by GitHub
parent ea4797bf38
commit 5acd3e66dd
7 changed files with 97 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rllib/agents/dqn/dqn_torch_policy.py
+2 -5
View File
@@ -52,7 +52,6 @@ class QLoss:
"mean_q": torch.mean(q_t_selected),
"min_q": torch.min(q_t_selected),
"max_q": torch.max(q_t_selected),
"td_error": self.td_error,
"mean_td_error": torch.mean(self.td_error),
}
@@ -250,10 +249,7 @@ def compute_q_values(policy, model, obs, explore, is_training=False):
def grad_process_and_td_error_fn(policy, optimizer, loss):
# Clip grads if configured.
info = apply_grad_clipping(policy, optimizer, loss)
# Add td-error to info dict.
info["td_error"] = policy.q_loss.td_error
return info
return apply_grad_clipping(policy, optimizer, loss)
def extra_action_out_fn(policy, input_dict, state_batches, model, action_dist):
@@ -270,6 +266,7 @@ DQNTorchPolicy = build_torch_policy(
postprocess_fn=postprocess_nstep_and_prio,
optimizer_fn=adam_optimizer,
extra_grad_process_fn=grad_process_and_td_error_fn,
extra_learn_fetches_fn=lambda policy: {"td_error": policy.q_loss.td_error},
extra_action_out_fn=extra_action_out_fn,
before_init=setup_early_mixins,
after_init=after_init,
rllib/agents/dqn/simple_q_torch_policy.py
+1 -1
View File
@@ -96,5 +96,5 @@ SimpleQTorchPolicy = build_torch_policy(
make_model_and_action_dist=build_q_model_and_distribution,
mixins=[TargetNetworkMixin],
action_distribution_fn=get_distribution_inputs_and_class,
stats_fn=lambda policy, config: {"td_error": policy.td_error},
extra_learn_fetches_fn=lambda policy: {"td_error": policy.td_error},
)
rllib/agents/es/tests/test_es.py
+5 -2
View File
@@ -9,14 +9,17 @@ from ray.rllib.utils.test_utils import check_compute_single_action, \
class TestES(unittest.TestCase):
def test_es_compilation(self):
"""Test whether an ESTrainer can be built on all frameworks."""
ray.init()
ray.init(num_cpus=2)
config = es.DEFAULT_CONFIG.copy()
# Keep it simple.
config["model"]["fcnet_hiddens"] = [10]
config["model"]["fcnet_activation"] = None
config["noise_size"] = 2500000
config["num_workers"] = 1
config["episodes_per_batch"] = 10
config["train_batch_size"] = 100
num_iterations = 2
num_iterations = 1
for _ in framework_iterator(config):
plain_config = config.copy()
rllib/agents/impala/impala.py
+2
View File
@@ -194,6 +194,8 @@ class BroadcastUpdateLearnerWeights:
metrics = _get_shared_metrics()
metrics.counters["num_weight_broadcasts"] += 1
actor.set_weights.remote(self.weights, _get_global_vars())
# Also update global vars of the local worker.
self.workers.local_worker().set_global_vars(_get_global_vars())
def record_steps_trained(item):
rllib/agents/impala/tests/test_impala.py
+19
View File
@@ -18,6 +18,25 @@ class TestIMPALA(unittest.TestCase):
def tearDownClass(cls) -> None:
ray.shutdown()
def test_impala_lr_schedule(self):
config = impala.DEFAULT_CONFIG.copy()
config["lr_schedule"] = [
[0, 0.0005],
[10000, 0.000001],
]
local_cfg = config.copy()
trainer = impala.ImpalaTrainer(config=local_cfg, env="CartPole-v0")
def get_lr(result):
return result["info"]["learner"]["default_policy"]["cur_lr"]
try:
r1 = trainer.train()
r2 = trainer.train()
assert get_lr(r2) < get_lr(r1), (r1, r2)
finally:
trainer.stop()
def test_impala_compilation(self):
"""Test whether an ImpalaTrainer can be built with both frameworks."""
config = impala.DEFAULT_CONFIG.copy()
rllib/policy/torch_policy.py
+54 -45
View File
@@ -40,23 +40,22 @@ class TorchPolicy(Policy):
"""
@DeveloperAPI
def __init__(self,
observation_space: gym.spaces.Space,
action_space: gym.spaces.Space,
config: TrainerConfigDict,
*,
model: ModelV2,
loss: Callable[
[Policy, ModelV2, type, SampleBatch], TensorType],
action_distribution_class: TorchDistributionWrapper,
action_sampler_fn: Callable[
[TensorType, List[TensorType]], Tuple[
TensorType, TensorType]] = None,
action_distribution_fn: Optional[Callable[
[Policy, ModelV2, TensorType, TensorType, TensorType],
Tuple[TensorType, type, List[TensorType]]]] = None,
max_seq_len: int = 20,
get_batch_divisibility_req: Optional[int] = None):
def __init__(
self,
observation_space: gym.spaces.Space,
action_space: gym.spaces.Space,
config: TrainerConfigDict,
*,
model: ModelV2,
loss: Callable[[Policy, ModelV2, type, SampleBatch], TensorType],
action_distribution_class: TorchDistributionWrapper,
action_sampler_fn: Callable[[TensorType, List[TensorType]], Tuple[
TensorType, TensorType]] = None,
action_distribution_fn: Optional[Callable[[
Policy, ModelV2, TensorType, TensorType, TensorType
], Tuple[TensorType, type, List[TensorType]]]] = None,
max_seq_len: int = 20,
get_batch_divisibility_req: Optional[int] = None):
"""Build a policy from policy and loss torch modules.
Note that model will be placed on GPU device if CUDA_VISIBLE_DEVICES
@@ -165,8 +164,8 @@ class TorchPolicy(Policy):
extra_fetches[SampleBatch.ACTION_PROB] = np.exp(logp)
extra_fetches[SampleBatch.ACTION_LOGP] = logp
return convert_to_non_torch_type(
(actions, state_out, extra_fetches))
return convert_to_non_torch_type((actions, state_out,
extra_fetches))
@override(Policy)
def compute_actions_from_trajectories(
@@ -183,8 +182,9 @@ class TorchPolicy(Policy):
with torch.no_grad():
# Create a view and pass that to Model as `input_dict`.
input_dict = self._lazy_tensor_dict(get_trajectory_view(
self.model, trajectories, is_training=False))
input_dict = self._lazy_tensor_dict(
get_trajectory_view(
self.model, trajectories, is_training=False))
# TODO: (sven) support RNNs w/ fast sampling.
state_batches = []
seq_lens = None
@@ -232,8 +232,8 @@ class TorchPolicy(Policy):
is_training=False)
else:
dist_class = self.dist_class
dist_inputs, state_out = self.model(
input_dict, state_batches, seq_lens)
dist_inputs, state_out = self.model(input_dict, state_batches,
seq_lens)
if not (isinstance(dist_class, functools.partial)
or issubclass(dist_class, TorchDistributionWrapper)):
@@ -270,10 +270,10 @@ class TorchPolicy(Policy):
actions: Union[List[TensorType], TensorType],
obs_batch: Union[List[TensorType], TensorType],
state_batches: Optional[List[TensorType]] = None,
prev_action_batch: Optional[
Union[List[TensorType], TensorType]] = None,
prev_reward_batch: Optional[
Union[List[TensorType], TensorType]] = None) -> TensorType:
prev_action_batch: Optional[Union[List[TensorType],
TensorType]] = None,
prev_reward_batch: Optional[Union[List[
TensorType], TensorType]] = None) -> TensorType:
if self.action_sampler_fn and self.action_distribution_fn is None:
raise ValueError("Cannot compute log-prob/likelihood w/o an "
@@ -314,8 +314,8 @@ class TorchPolicy(Policy):
@override(Policy)
@DeveloperAPI
def learn_on_batch(self, postprocessed_batch: SampleBatch) -> Dict[
str, TensorType]:
def learn_on_batch(
self, postprocessed_batch: SampleBatch) -> Dict[str, TensorType]:
# Get batch ready for RNNs, if applicable.
pad_batch_to_sequences_of_same_size(
postprocessed_batch,
@@ -331,6 +331,7 @@ class TorchPolicy(Policy):
loss_out = self.model.custom_loss(loss_out, train_batch)
assert len(loss_out) == len(self._optimizers)
# assert not any(torch.isnan(l) for l in loss_out)
fetches = self.extra_compute_grad_fetches()
# Loop through all optimizers.
grad_info = {"allreduce_latency": 0.0}
@@ -370,7 +371,7 @@ class TorchPolicy(Policy):
grad_info["allreduce_latency"] /= len(self._optimizers)
grad_info.update(self.extra_grad_info(train_batch))
return {LEARNER_STATS_KEY: grad_info}
return dict(fetches, **{LEARNER_STATS_KEY: grad_info})
@override(Policy)
@DeveloperAPI
@@ -380,6 +381,7 @@ class TorchPolicy(Policy):
loss_out = force_list(
self._loss(self, self.model, self.dist_class, train_batch))
assert len(loss_out) == len(self._optimizers)
fetches = self.extra_compute_grad_fetches()
grad_process_info = {}
grads = []
@@ -399,7 +401,7 @@ class TorchPolicy(Policy):
grad_info = self.extra_grad_info(train_batch)
grad_info.update(grad_process_info)
return grads, {LEARNER_STATS_KEY: grad_info}
return grads, dict(fetches, **{LEARNER_STATS_KEY: grad_info})
@override(Policy)
@DeveloperAPI
@@ -466,10 +468,8 @@ class TorchPolicy(Policy):
super().set_state(state)
@DeveloperAPI
def extra_grad_process(
self,
optimizer: "torch.optim.Optimizer",
loss: TensorType):
def extra_grad_process(self, optimizer: "torch.optim.Optimizer",
loss: TensorType):
"""Called after each optimizer.zero_grad() + loss.backward() call.
Called for each self._optimizers/loss-value pair.
@@ -486,12 +486,20 @@ class TorchPolicy(Policy):
"""
return {}
@DeveloperAPI
def extra_compute_grad_fetches(self) -> Dict[str, any]:
"""Extra values to fetch and return from compute_gradients().
Returns:
Dict[str, any]: Extra fetch dict to be added to the fetch dict
of the compute_gradients call.
"""
return {LEARNER_STATS_KEY: {}} # e.g, stats, td error, etc.
@DeveloperAPI
def extra_action_out(
self,
input_dict: Dict[str, TensorType],
state_batches: List[TensorType],
model: TorchModelV2,
self, input_dict: Dict[str, TensorType],
state_batches: List[TensorType], model: TorchModelV2,
action_dist: TorchDistributionWrapper) -> Dict[str, TensorType]:
"""Returns dict of extra info to include in experience batch.
@@ -509,8 +517,8 @@ class TorchPolicy(Policy):
return {}
@DeveloperAPI
def extra_grad_info(self, train_batch: SampleBatch) -> Dict[
str, TensorType]:
def extra_grad_info(self,
train_batch: SampleBatch) -> Dict[str, TensorType]:
"""Return dict of extra grad info.
Args:
@@ -524,8 +532,9 @@ class TorchPolicy(Policy):
return {}
@DeveloperAPI
def optimizer(self) -> Union[
List["torch.optim.Optimizer"], "torch.optim.Optimizer"]:
def optimizer(
self
) -> Union[List["torch.optim.Optimizer"], "torch.optim.Optimizer"]:
"""Custom the local PyTorch optimizer(s) to use.
Returns:
@@ -560,8 +569,8 @@ class TorchPolicy(Policy):
def _lazy_tensor_dict(self, postprocessed_batch):
train_batch = UsageTrackingDict(postprocessed_batch)
train_batch.set_get_interceptor(functools.partial(
convert_to_torch_tensor, device=self.device))
train_batch.set_get_interceptor(
functools.partial(convert_to_torch_tensor, device=self.device))
return train_batch
rllib/policy/torch_policy_template.py
+14 -1
View File
@@ -1,4 +1,4 @@
from ray.rllib.policy.policy import Policy
from ray.rllib.policy.policy import Policy, LEARNER_STATS_KEY
from ray.rllib.policy.torch_policy import TorchPolicy
from ray.rllib.models.catalog import ModelCatalog
from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
@@ -19,6 +19,7 @@ def build_torch_policy(name,
postprocess_fn=None,
extra_action_out_fn=None,
extra_grad_process_fn=None,
extra_learn_fetches_fn=None,
optimizer_fn=None,
validate_spaces=None,
before_init=None,
@@ -47,6 +48,8 @@ def build_torch_policy(name,
returns a dict of extra values to include in experiences.
extra_grad_process_fn (Optional[callable]): Optional callable that is
called after gradients are computed and returns processing info.
extra_learn_fetches_fn (func): optional function that returns a dict of
extra values to fetch from the policy after loss evaluation.
optimizer_fn (Optional[callable]): Optional callable that returns a
torch optimizer given the policy and config.
validate_spaces (Optional[callable]): Optional callable that takes the
@@ -179,6 +182,16 @@ def build_torch_policy(name,
else:
return TorchPolicy.extra_grad_process(self, optimizer, loss)
@override(TorchPolicy)
def extra_compute_grad_fetches(self):
if extra_learn_fetches_fn:
fetches = convert_to_non_torch_type(
extra_learn_fetches_fn(self))
# Auto-add empty learner stats dict if needed.
return dict({LEARNER_STATS_KEY: {}}, **fetches)
else:
return TorchPolicy.extra_compute_grad_fetches(self)
@override(TorchPolicy)
def apply_gradients(self, gradients):
if apply_gradients_fn: