[rllib] Make timestep a required arg for exploration classes (#7380)

This commit is contained in:
Eric Liang
2020-03-04 13:00:37 -08:00
committed by GitHub
parent fddeb6809c
commit 596b39e36a
14 changed files with 98 additions and 155 deletions
+2 -2
View File
@@ -122,8 +122,8 @@ class DDPGTFPolicy(DDPGPostprocessing, TFPolicy):
# Action outputs
with tf.variable_scope(ACTION_SCOPE):
self.output_actions, _ = self.exploration.get_exploration_action(
policy_out, Deterministic, self.policy_model, explore,
timestep)
policy_out, Deterministic, self.policy_model, timestep,
explore)
# Replay inputs
self.obs_t = tf.placeholder(
+1 -1
View File
@@ -224,7 +224,7 @@ def sample_action_from_q_network(policy, q_model, input_dict, obs_space,
policy.output_actions, policy.sampled_action_logp = \
policy.exploration.get_exploration_action(
policy.q_values, Categorical, q_model, explore, timestep)
policy.q_values, Categorical, q_model, timestep, explore)
# Noise vars for Q network except for layer normalization vars.
if config["parameter_noise"]:
+1 -1
View File
@@ -111,7 +111,7 @@ def simple_sample_action_from_q_network(policy, q_model, input_dict, obs_space,
policy.output_actions, policy.sampled_action_logp = \
policy.exploration.get_exploration_action(
policy.q_values, Categorical, q_model, explore, timestep)
policy.q_values, Categorical, q_model, timestep, explore)
return policy.output_actions, policy.sampled_action_logp
+1 -1
View File
@@ -116,7 +116,7 @@ def build_action_output(policy, model, input_dict, obs_space, action_space,
policy.output_actions, policy.sampled_action_logp = \
policy.exploration.get_exploration_action(
distribution_inputs, action_dist_class, model, explore, timestep)
distribution_inputs, action_dist_class, model, timestep, explore)
return policy.output_actions, policy.sampled_action_logp
+2 -2
View File
@@ -184,8 +184,8 @@ class DynamicTFPolicy(TFPolicy):
model_out,
self.dist_class,
self.model,
explore=explore,
timestep=timestep)
timestep,
explore=explore)
# Phase 1 init.
sess = tf.get_default_session() or tf.Session()
+11 -4
View File
@@ -332,8 +332,15 @@ def build_eager_tf_policy(name,
if action_sampler_fn is not None:
state_out = []
action, logp = action_sampler_fn(
self, self.model, input_dict, self.observation_space,
self.action_space, explore, self.config, timestep)
self,
self.model,
input_dict,
self.observation_space,
self.action_space,
explore,
self.config,
timestep=timestep
if timestep is not None else self.global_timestep)
# Use Exploration object.
else:
with tf.variable_creator_scope(_disallow_var_creation):
@@ -343,9 +350,9 @@ def build_eager_tf_policy(name,
model_out,
self.dist_class,
self.model,
explore=explore,
timestep=timestep
if timestep is not None else self.global_timestep)
if timestep is not None else self.global_timestep,
explore=explore)
extra_fetches = {}
if logp is not None:
+2 -2
View File
@@ -86,9 +86,9 @@ class TorchPolicy(Policy):
action_dist = None
actions, logp = \
self.exploration.get_exploration_action(
logits, self.dist_class, self.model, explore,
logits, self.dist_class, self.model,
timestep if timestep is not None else
self.global_timestep)
self.global_timestep, explore)
input_dict[SampleBatch.ACTIONS] = actions
extra_action_out = self.extra_action_out(input_dict, state_batches,
+14 -47
View File
@@ -1,10 +1,11 @@
import numpy as np
from typing import Union
from ray.rllib.utils.annotations import override
from ray.rllib.utils.exploration.exploration import Exploration
from ray.rllib.utils.exploration.exploration import Exploration, TensorType
from ray.rllib.utils.framework import try_import_tf, try_import_torch, \
get_variable
from ray.rllib.utils.schedules import PiecewiseSchedule
from ray.rllib.models.modelv2 import ModelV2
tf = try_import_tf()
torch, _ = try_import_torch()
@@ -26,7 +27,7 @@ class EpsilonGreedy(Exploration):
epsilon_schedule=None,
framework="tf",
**kwargs):
"""
"""Create an EpsilonGreedy exploration class.
Args:
action_space (Space): The gym action space used by the environment.
@@ -54,11 +55,11 @@ class EpsilonGreedy(Exploration):
@override(Exploration)
def get_exploration_action(self,
distribution_inputs,
action_dist_class=None,
model=None,
explore=True,
timestep=None):
distribution_inputs: TensorType,
action_dist_class: type,
model: ModelV2,
timestep: Union[int, TensorType],
explore: bool = True):
if self.framework == "tf":
return self._get_tf_exploration_action_op(distribution_inputs,
@@ -68,7 +69,7 @@ class EpsilonGreedy(Exploration):
explore, timestep)
def _get_tf_exploration_action_op(self, q_values, explore, timestep):
"""Tf method to produce the tf op for an epsilon exploration action.
"""TF method to produce the tf op for an epsilon exploration action.
Args:
q_values (Tensor): The Q-values coming from some q-model.
@@ -104,10 +105,7 @@ class EpsilonGreedy(Exploration):
),
false_fn=lambda: exploit_action)
# Increment `last_timestep` by 1 (or set to `timestep`).
assign_op = \
tf.assign_add(self.last_timestep, 1) if timestep is None else \
tf.assign(self.last_timestep, timestep)
assign_op = tf.assign(self.last_timestep, timestep)
with tf.control_dependencies([assign_op]):
return action, tf.zeros_like(action, dtype=tf.float32)
@@ -120,10 +118,7 @@ class EpsilonGreedy(Exploration):
Returns:
torch.Tensor: The exploration-action.
"""
# Set last timestep or (if not given) increase by one.
self.last_timestep = timestep if timestep is not None else \
self.last_timestep + 1
self.last_timestep = timestep
_, exploit_action = torch.max(q_values, 1)
action_logp = torch.zeros_like(exploit_action)
@@ -153,33 +148,5 @@ class EpsilonGreedy(Exploration):
@override(Exploration)
def get_info(self):
"""Returns the current epsilon value.
Returns:
Union[float,tf.Tensor[float]]: The current epsilon value.
"""
return self.epsilon_schedule(self.last_timestep)
@override(Exploration)
def get_state(self):
return [self.last_timestep]
@override(Exploration)
def set_state(self, state):
if self.framework == "tf" and tf.executing_eagerly() is False:
update_op = tf.assign(self.last_timestep, state)
with tf.control_dependencies([update_op]):
return tf.no_op()
self.last_timestep = state
@override(Exploration)
def reset_state(self):
return self.set_state(0)
@classmethod
@override(Exploration)
def merge_states(cls, exploration_objects):
timesteps = [e.get_state() for e in exploration_objects]
if exploration_objects[0].framework == "tf":
return tf.reduce_sum(timesteps)
return np.sum(timesteps)
eps = self.epsilon_schedule(self.last_timestep)
return {"cur_epsilon": eps}
+30 -75
View File
@@ -1,4 +1,8 @@
from ray.rllib.utils.framework import check_framework, try_import_tf
from gym.spaces import Space
from ray.rllib.utils.framework import check_framework, try_import_tf, \
TensorType
from ray.rllib.models.modelv2 import ModelV2
from typing import Union
tf = try_import_tf()
@@ -12,18 +16,15 @@ class Exploration:
"""
def __init__(self,
action_space=None,
*,
num_workers=None,
worker_index=None,
framework="tf"):
action_space: Space,
num_workers: int = 0,
worker_index: int = 0,
framework: str = "tf"):
"""
Args:
action_space (Optional[gym.spaces.Space]): The action space in
which to explore.
num_workers (Optional[int]): The overall number of workers used.
worker_index (Optional[int]): The index of the Worker using this
Exploration.
action_space (Space): The action space in which to explore.
num_workers (int): The overall number of workers used.
worker_index (int): The index of the worker using this class.
framework (str): One of "tf" or "torch".
"""
self.action_space = action_space
@@ -32,30 +33,28 @@ class Exploration:
self.framework = check_framework(framework)
def get_exploration_action(self,
distribution_inputs,
action_dist_class,
model=None,
explore=True,
timestep=None):
distribution_inputs: TensorType,
action_dist_class: type,
model: ModelV2,
timestep: Union[int, TensorType],
explore: bool = True):
"""Returns a (possibly) exploratory action and its log-likelihood.
Given the Model's logits outputs and action distribution, returns an
exploratory action.
Args:
distribution_inputs (any): The output coming from the model,
distribution_inputs (TensorType): The output coming from the model,
ready for parameterizing a distribution
(e.g. q-values or PG-logits).
action_dist_class (class): The action distribution class
to use.
model (ModelV2): The Model object.
timestep (int|TensorType): The current sampling time step. It can
be a tensor for TF graph mode, otherwise an integer.
explore (bool): True: "Normal" exploration behavior.
False: Suppress all exploratory behavior and return
a deterministic action.
timestep (int): The current sampling time step. If None, the
component should try to use an internal counter, which it
then increments by 1. If provided, will set the internal
counter to the given value.
Returns:
Tuple:
@@ -66,21 +65,21 @@ class Exploration:
pass
def get_loss_exploration_term(self,
model_output,
model=None,
action_dist=None,
action_sample=None):
model_output: TensorType,
model: ModelV2,
action_dist: type,
action_sample: TensorType = None):
"""Returns an extra loss term to be added to a loss.
Args:
model_output (any): The Model's output Tensor(s).
model_output (TensorType): The Model's output Tensor(s).
model (ModelV2): The Model object.
action_dist: The ActionDistribution object resulting from
`model_output`. TODO: Or the class?
action_sample (any): An optional action sample.
action_sample (TensorType): An optional action sample.
Returns:
any: The extra loss term to add to the loss.
TensorType: The extra loss term to add to the loss.
"""
pass # TODO(sven): implement for some example Exploration class.
@@ -91,51 +90,7 @@ class Exploration:
set_state!), but rather useful (e.g. debugging) information.
Returns:
any: A description of the Exploration (not necessarily its state).
dict: A description of the Exploration (not necessarily its state).
This may include tf.ops as values in graph mode.
"""
if self.framework == "tf":
return tf.no_op()
def get_state(self):
"""Returns the current exploration state.
Returns:
List[any]: The current state (or a tf-op thereof).
"""
return []
def set_state(self, state):
"""Sets the current state of the Exploration to the given value.
Or returns a tf op that will do the set.
Args:
state (List[any]): The new state to set.
Returns:
Union[None,tf.op]: If framework=tf, the op that handles the update.
"""
pass
def reset_state(self):
"""Resets the exploration's state.
Returns:
Union[None,tf.op]: If framework=tf, the op that handles the reset.
"""
pass
@classmethod
def merge_states(cls, exploration_objects):
"""Returns the merged states of all exploration_objects as a value.
Or a tf.Tensor (whose execution will trigger the merge).
Args:
exploration_objects (List[Exploration]): All Exploration objects,
whose states have to be merged somehow.
Returns:
The merged value or a tf.op to execute.
"""
pass
return {}
+10 -7
View File
@@ -1,9 +1,12 @@
from typing import Union
from ray.rllib.utils.annotations import override
from ray.rllib.utils.exploration.exploration import Exploration
from ray.rllib.utils.exploration.random import Random
from ray.rllib.utils.framework import try_import_tf, try_import_torch, \
get_variable
get_variable, TensorType
from ray.rllib.utils.schedules.piecewise_schedule import PiecewiseSchedule
from ray.rllib.models.modelv2 import ModelV2
tf = try_import_tf()
torch, _ = try_import_torch()
@@ -54,7 +57,7 @@ class GaussianNoise(Exploration):
self.random_timesteps = random_timesteps
self.random_exploration = Random(
action_space, framework=self.framework)
action_space, framework=self.framework, **kwargs)
self.stddev = stddev
# The `scale` annealing schedule.
self.scale_schedule = scale_schedule or PiecewiseSchedule(
@@ -69,11 +72,11 @@ class GaussianNoise(Exploration):
@override(Exploration)
def get_exploration_action(self,
distribution_inputs,
action_dist_class,
model=None,
explore=True,
timestep=None):
distribution_inputs: TensorType,
action_dist_class: type,
model: ModelV2,
timestep: Union[int, TensorType],
explore: bool = True):
# Adds IID Gaussian noise for exploration, TD3-style.
action_dist = action_dist_class(distribution_inputs, model)
@@ -17,7 +17,8 @@ class PerWorkerEpsilonGreedy(EpsilonGreedy):
worker_index=0,
framework="tf",
**kwargs):
"""
"""Create a PerWorkerEpsilonGreedy exploration class.
Args:
action_space (Space): The gym action space used by the environment.
num_workers (Optional[int]): The overall number of workers used.
+8 -6
View File
@@ -1,10 +1,12 @@
from gym.spaces import Discrete, MultiDiscrete, Tuple
from typing import Union
from ray.rllib.utils.annotations import override
from ray.rllib.utils.exploration.exploration import Exploration
from ray.rllib.utils.framework import try_import_tf, try_import_torch, \
tf_function
tf_function, TensorType
from ray.rllib.utils.tuple_actions import TupleActions
from ray.rllib.models.modelv2 import ModelV2
tf = try_import_tf()
torch, _ = try_import_torch()
@@ -38,11 +40,11 @@ class Random(Exploration):
@override(Exploration)
def get_exploration_action(self,
distribution_inputs,
action_dist_class,
model=None,
explore=True,
timestep=None):
distribution_inputs: TensorType,
action_dist_class: type,
model: ModelV2,
timestep: Union[int, TensorType],
explore: bool = True):
# Instantiate the distribution object.
action_dist = action_dist_class(distribution_inputs, model)
if self.framework == "tf":
+10 -6
View File
@@ -1,7 +1,11 @@
from typing import Union
from ray.rllib.utils.annotations import override
from ray.rllib.utils.exploration.exploration import Exploration
from ray.rllib.utils.framework import try_import_tf, try_import_torch
from ray.rllib.utils.framework import try_import_tf, try_import_torch, \
TensorType
from ray.rllib.utils.tuple_actions import TupleActions
from ray.rllib.models.modelv2 import ModelV2
tf = try_import_tf()
torch, _ = try_import_torch()
@@ -45,11 +49,11 @@ class StochasticSampling(Exploration):
@override(Exploration)
def get_exploration_action(self,
distribution_inputs,
action_dist_class,
model=None,
explore=True,
timestep=None):
distribution_inputs: TensorType,
action_dist_class: type,
model: ModelV2,
timestep: Union[int, TensorType],
explore: bool = True):
kwargs = self.static_params.copy()
# TODO(sven): create schedules for these via easy-config patterns
+4
View File
@@ -1,8 +1,12 @@
import logging
import os
from typing import Any
logger = logging.getLogger(__name__)
# Represents a generic tensor type.
TensorType = Any
def check_framework(framework="tf"):
"""