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

[RLlib] Redo: Make TFModelV2 fully modular like TorchModelV2 (soft-deprecate register_variables, unify var names wrt torch). (#13363)

Browse Source
This commit is contained in:
Sven Mika
2021-01-14 14:44:33 +01:00
committed by GitHub
parent 33b092de28
commit 56878221ed
33 changed files with 259 additions and 153 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rllib/BUILD
+9 -2
View File
@@ -1089,6 +1089,13 @@ py_test(
srcs = ["models/tests/test_distributions.py"]
)
py_test(
name = "test_models",
tags = ["models"],
size = "small",
srcs = ["models/tests/test_models.py"]
)
py_test(
name = "test_preprocessors",
tags = ["models"],
@@ -1541,7 +1548,7 @@ py_test(
name = "examples/batch_norm_model_dqn_tf",
main = "examples/batch_norm_model.py",
tags = ["examples", "examples_B"],
size = "large", # DQN learns much slower with BatchNorm.
size = "medium",
srcs = ["examples/batch_norm_model.py"],
args = ["--as-test", "--run=DQN", "--stop-reward=70"]
)
@@ -1559,7 +1566,7 @@ py_test(
name = "examples/batch_norm_model_ddpg_tf",
main = "examples/batch_norm_model.py",
tags = ["examples", "examples_B"],
size = "medium",
size = "small",
srcs = ["examples/batch_norm_model.py"],
args = ["--run=DDPG", "--stop-iters=1"]
)
rllib/agents/ddpg/ddpg_tf_model.py
-3
View File
@@ -89,7 +89,6 @@ class DDPGTFModel(TFModelV2):
actor_out = tf.keras.layers.Lambda(lambda_)(actor_out)
self.policy_model = tf.keras.Model(self.model_out, actor_out)
self.register_variables(self.policy_model.variables)
# Build the Q-model(s).
self.actions_input = tf.keras.layers.Input(
@@ -116,12 +115,10 @@ class DDPGTFModel(TFModelV2):
return q_net
self.q_model = build_q_net("q", self.model_out, self.actions_input)
self.register_variables(self.q_model.variables)
if twin_q:
self.twin_q_model = build_q_net("twin_q", self.model_out,
self.actions_input)
self.register_variables(self.twin_q_model.variables)
else:
self.twin_q_model = None
rllib/agents/dqn/distributional_q_tf_model.py
-2
View File
@@ -162,13 +162,11 @@ class DistributionalQTFModel(TFModelV2):
q_out = build_action_value(name + "/action_value/", self.model_out)
self.q_value_head = tf.keras.Model(self.model_out, q_out)
self.register_variables(self.q_value_head.variables)
if dueling:
state_out = build_state_score(name + "/state_value/",
self.model_out)
self.state_value_head = tf.keras.Model(self.model_out, state_out)
self.register_variables(self.state_value_head.variables)
def get_q_value_distributions(self,
model_out: TensorType) -> List[TensorType]:
rllib/agents/sac/sac_tf_model.py
-6
View File
@@ -91,8 +91,6 @@ class SACTFModel(TFModelV2):
])
self.shift_and_log_scale_diag = self.action_model(self.model_out)
self.register_variables(self.action_model.variables)
self.actions_input = None
if not self.discrete:
self.actions_input = tf.keras.layers.Input(
@@ -123,12 +121,10 @@ class SACTFModel(TFModelV2):
return q_net
self.q_net = build_q_net("q", self.model_out, self.actions_input)
self.register_variables(self.q_net.variables)
if twin_q:
self.twin_q_net = build_q_net("twin_q", self.model_out,
self.actions_input)
self.register_variables(self.twin_q_net.variables)
else:
self.twin_q_net = None
@@ -147,8 +143,6 @@ class SACTFModel(TFModelV2):
target_entropy = -np.prod(action_space.shape)
self.target_entropy = target_entropy
self.register_variables([self.log_alpha])
def get_q_values(self,
model_out: TensorType,
actions: Optional[TensorType] = None) -> TensorType:
rllib/examples/batch_norm_model.py
+6 -4
View File
@@ -6,7 +6,7 @@ import os
import ray
from ray import tune
from ray.rllib.examples.models.batch_norm_model import BatchNormModel, \
TorchBatchNormModel
KerasBatchNormModel, TorchBatchNormModel
from ray.rllib.models import ModelCatalog
from ray.rllib.utils.framework import try_import_tf
from ray.rllib.utils.test_utils import check_learning_achieved
@@ -26,13 +26,15 @@ if __name__ == "__main__":
ray.init()
ModelCatalog.register_custom_model(
"bn_model", TorchBatchNormModel if args.torch else BatchNormModel)
"bn_model", TorchBatchNormModel if args.torch else KerasBatchNormModel
if args.run != "PPO" else BatchNormModel)
config = {
"env": "Pendulum-v0" if args.run == "DDPG" else "CartPole-v0",
"env": "Pendulum-v0" if args.run in ["DDPG", "SAC"] else "CartPole-v0",
"model": {
"custom_model": "bn_model",
},
"lr": 0.0003,
# Use GPUs iff `RLLIB_NUM_GPUS` env var set to > 0.
"num_gpus": int(os.environ.get("RLLIB_NUM_GPUS", "0")),
"num_workers": 0,
@@ -45,7 +47,7 @@ if __name__ == "__main__":
"episode_reward_mean": args.stop_reward,
}
results = tune.run(args.run, stop=stop, config=config, verbose=1)
results = tune.run(args.run, stop=stop, config=config, verbose=2)
if args.as_test:
check_learning_achieved(results, args.stop_reward)
rllib/examples/custom_env.py
-1
View File
@@ -71,7 +71,6 @@ class CustomModel(TFModelV2):
model_config, name)
self.model = FullyConnectedNetwork(obs_space, action_space,
num_outputs, model_config, name)
self.register_variables(self.model.variables())
def forward(self, input_dict, state, seq_lens):
return self.model.forward(input_dict, state, seq_lens)
rllib/examples/custom_keras_model.py
-2
View File
@@ -48,7 +48,6 @@ class MyKerasModel(TFModelV2):
activation=None,
kernel_initializer=normc_initializer(0.01))(layer_1)
self.base_model = tf.keras.Model(self.inputs, [layer_out, value_out])
self.register_variables(self.base_model.variables)
def forward(self, input_dict, state, seq_lens):
model_out, self._value_out = self.base_model(input_dict["obs"])
@@ -84,7 +83,6 @@ class MyKerasQModel(DistributionalQTFModel):
activation=tf.nn.relu,
kernel_initializer=normc_initializer(1.0))(layer_1)
self.base_model = tf.keras.Model(self.inputs, layer_out)
self.register_variables(self.base_model.variables)
# Implement the core forward method
def forward(self, input_dict, state, seq_lens):
rllib/examples/models/autoregressive_action_model.py
-2
View File
@@ -69,14 +69,12 @@ class AutoregressiveActionModel(TFModelV2):
# Base layers
self.base_model = tf.keras.Model(obs_input, [context, value_out])
self.register_variables(self.base_model.variables)
self.base_model.summary()
# Autoregressive action sampler
self.action_model = tf.keras.Model([ctx_input, a1_input],
[a1_logits, a2_logits])
self.action_model.summary()
self.register_variables(self.action_model.variables)
def forward(self, input_dict, state, seq_lens):
context, self._value_out = self.base_model(input_dict["obs"])
rllib/examples/models/batch_norm_model.py
+77 -61
View File
@@ -13,13 +13,73 @@ tf1, tf, tfv = try_import_tf()
torch, nn = try_import_torch()
class KerasBatchNormModel(TFModelV2):
"""Keras version of above BatchNormModel with exactly the same structure.
IMORTANT NOTE: This model will not work with PPO due to a bug in keras
that surfaces when having more than one input placeholder (here: `inputs`
and `is_training`) AND using the `make_tf_callable` helper (e.g. used by
PPO), in which auto-placeholders are generated, then passed through the
tf.keras. models.Model. In this last step, the connection between 1) the
provided value in the auto-placeholder and 2) the keras `is_training`
Input is broken and keras complains.
Use the below `BatchNormModel` (a non-keras based TFModelV2), instead.
"""
def __init__(self, obs_space, action_space, num_outputs, model_config,
name):
super().__init__(obs_space, action_space, num_outputs, model_config,
name)
inputs = tf.keras.layers.Input(shape=obs_space.shape, name="inputs")
# Have to batch the is_training flag (its batch size will always be 1).
is_training = tf.keras.layers.Input(
shape=(), dtype=tf.bool, batch_size=1, name="is_training")
last_layer = inputs
hiddens = [256, 256]
for i, size in enumerate(hiddens):
label = "fc{}".format(i)
last_layer = tf.keras.layers.Dense(
units=size,
kernel_initializer=normc_initializer(1.0),
activation=tf.nn.tanh,
name=label)(last_layer)
# Add a batch norm layer
last_layer = tf.keras.layers.BatchNormalization()(
last_layer, training=is_training[0])
output = tf.keras.layers.Dense(
units=self.num_outputs,
kernel_initializer=normc_initializer(0.01),
activation=None,
name="fc_out")(last_layer)
value_out = tf.keras.layers.Dense(
units=1,
kernel_initializer=normc_initializer(0.01),
activation=None,
name="value_out")(last_layer)
self.base_model = tf.keras.models.Model(
inputs=[inputs, is_training], outputs=[output, value_out])
@override(ModelV2)
def forward(self, input_dict, state, seq_lens):
# Have to batch the is_training flag (B=1).
out, self._value_out = self.base_model(
[input_dict["obs"],
tf.expand_dims(input_dict["is_training"], 0)])
return out, []
@override(ModelV2)
def value_function(self):
return tf.reshape(self._value_out, [-1])
class BatchNormModel(TFModelV2):
"""Example of a TFModelV2 that is built w/o using tf.keras.
NOTE: This example does not work when using a keras-based TFModelV2 due
to a bug in keras related to missing values for input placeholders, even
NOTE: The above keras-based example model does not work with PPO (due to
a bug in keras related to missing values for input placeholders, even
though these input values have been provided in a forward pass through the
actual keras Model.
actual keras Model).
All Model logic (layers) is defined in the `forward` method (incl.
the batch_normalization layers). Also, all variables are registered
@@ -65,7 +125,21 @@ class BatchNormModel(TFModelV2):
kernel_initializer=normc_initializer(1.0),
activation=None,
name="vf")
# Register variables.
# NOTE: This is not the recommended way of doing things. We would
# prefer creating keras-style Layers like it's done in the
# `KerasBatchNormModel` class above and then have TFModelV2 auto-detect
# the created vars. However, since there is a bug
# in keras/tf that prevents us from using that KerasBatchNormModel
# example (see comments above), we do variable registration the old,
# manual way for this example Model here.
if not self._registered:
# Register already auto-detected variables (from the wrapping
# Model, e.g. DQNTFModel).
self.register_variables(self.variables())
# Then register everything we added to the graph in this `forward`
# call.
self.register_variables(
tf1.get_collection(
tf1.GraphKeys.TRAINABLE_VARIABLES, scope=".+/model/.+"))
@@ -78,64 +152,6 @@ class BatchNormModel(TFModelV2):
return tf.reshape(self._value_out, [-1])
class KerasBatchNormModel(TFModelV2):
"""Keras version of above BatchNormModel with exactly the same structure.
IMORTANT NOTE: This model will not work with PPO due to a bug in keras
that surfaces when having more than one input placeholder (here: `inputs`
and `is_training`) AND using the `make_tf_callable` helper (e.g. used by
PPO), in which auto-placeholders are generated, then passed through the
tf.keras. models.Model. In this last step, the connection between 1) the
provided value in the auto-placeholder and 2) the keras `is_training`
Input is broken and keras complains.
Use the above `BatchNormModel` (a non-keras based TFModelV2), instead.
"""
def __init__(self, obs_space, action_space, num_outputs, model_config,
name):
super().__init__(obs_space, action_space, num_outputs, model_config,
name)
inputs = tf.keras.layers.Input(shape=obs_space.shape, name="inputs")
is_training = tf.keras.layers.Input(
shape=(), dtype=tf.bool, batch_size=1, name="is_training")
last_layer = inputs
hiddens = [256, 256]
for i, size in enumerate(hiddens):
label = "fc{}".format(i)
last_layer = tf.keras.layers.Dense(
units=size,
kernel_initializer=normc_initializer(1.0),
activation=tf.nn.tanh,
name=label)(last_layer)
# Add a batch norm layer
last_layer = tf.keras.layers.BatchNormalization()(
last_layer, training=is_training[0])
output = tf.keras.layers.Dense(
units=self.num_outputs,
kernel_initializer=normc_initializer(0.01),
activation=None,
name="fc_out")(last_layer)
value_out = tf.keras.layers.Dense(
units=1,
kernel_initializer=normc_initializer(0.01),
activation=None,
name="value_out")(last_layer)
self.base_model = tf.keras.models.Model(
inputs=[inputs, is_training], outputs=[output, value_out])
self.register_variables(self.base_model.variables)
@override(ModelV2)
def forward(self, input_dict, state, seq_lens):
out, self._value_out = self.base_model(
[input_dict["obs"], input_dict["is_training"]])
return out, []
@override(ModelV2)
def value_function(self):
return tf.reshape(self._value_out, [-1])
class TorchBatchNormModel(TorchModelV2, nn.Module):
"""Example of a TorchModelV2 using batch normalization."""
capture_index = 0
rllib/examples/models/centralized_critic_models.py
-4
View File
@@ -23,7 +23,6 @@ class CentralizedCriticModel(TFModelV2):
# Base of the model
self.model = FullyConnectedNetwork(obs_space, action_space,
num_outputs, model_config, name)
self.register_variables(self.model.variables())
# Central VF maps (obs, opp_obs, opp_act) -> vf_pred
obs = tf.keras.layers.Input(shape=(6, ), name="obs")
@@ -37,7 +36,6 @@ class CentralizedCriticModel(TFModelV2):
1, activation=None, name="c_vf_out")(central_vf_dense)
self.central_vf = tf.keras.Model(
inputs=[obs, opp_obs, opp_act], outputs=central_vf_out)
self.register_variables(self.central_vf.variables)
@override(ModelV2)
def forward(self, input_dict, state, seq_lens):
@@ -79,11 +77,9 @@ class YetAnotherCentralizedCriticModel(TFModelV2):
num_outputs,
model_config,
name + "_action")
self.register_variables(self.action_model.variables())
self.value_model = FullyConnectedNetwork(obs_space, action_space, 1,
model_config, name + "_vf")
self.register_variables(self.value_model.variables())
def forward(self, input_dict, state, seq_lens):
self._value_out, _ = self.value_model({
rllib/examples/models/cnn_plus_fc_concat_model.py
-2
View File
@@ -53,7 +53,6 @@ class CNNPlusFCConcatModel(TFModelV2):
name="cnn_{}".format(i))
concat_size += cnn.num_outputs
self.cnns[i] = cnn
self.register_variables(cnn.variables())
# Discrete inputs -> One-hot encode.
elif isinstance(component, Discrete):
concat_size += component.n
@@ -82,7 +81,6 @@ class CNNPlusFCConcatModel(TFModelV2):
kernel_initializer=normc_initializer(0.01))(concat_layer)
self.logits_and_value_model = tf.keras.models.Model(
concat_layer, [logits_layer, value_layer])
self.register_variables(self.logits_and_value_model.variables)
else:
self.num_outputs = concat_size
rllib/examples/models/custom_loss_model.py
-1
View File
@@ -29,7 +29,6 @@ class CustomLossModel(TFModelV2):
num_outputs,
model_config,
name="fcnet")
self.register_variables(self.fcnet.variables())
@override(ModelV2)
def forward(self, input_dict, state, seq_lens):
rllib/examples/models/eager_model.py
-1
View File
@@ -40,7 +40,6 @@ class EagerModel(TFModelV2):
out = tf.keras.layers.Lambda(lambda_)(out)
self.base_model = tf.keras.models.Model(inputs, [out, value_out])
self.register_variables(self.base_model.variables)
@override(ModelV2)
def forward(self, input_dict, state, seq_lens):
rllib/examples/models/mobilenet_v2_with_lstm_models.py
-1
View File
@@ -68,7 +68,6 @@ class MobileV2PlusRNNModel(RecurrentNetwork):
self.rnn_model = tf.keras.Model(
inputs=[inputs, seq_in, state_in_h, state_in_c],
outputs=[logits, values, state_h, state_c])
self.register_variables(self.rnn_model.variables)
self.rnn_model.summary()
@override(RecurrentNetwork)
rllib/examples/models/parametric_actions_model.py
-1
View File
@@ -36,7 +36,6 @@ class ParametricActionsModel(DistributionalQTFModel):
self.action_embed_model = FullyConnectedNetwork(
Box(-1, 1, shape=true_obs_shape), action_space, action_embed_size,
model_config, name + "_action_embed")
self.register_variables(self.action_embed_model.variables())
def forward(self, input_dict, state, seq_lens):
# Extract the available actions tensor from the observation.
rllib/examples/models/rnn_model.py
-1
View File
@@ -54,7 +54,6 @@ class RNNModel(RecurrentNetwork):
self.rnn_model = tf.keras.Model(
inputs=[input_layer, seq_in, state_in_h, state_in_c],
outputs=[logits, values, state_h, state_c])
self.register_variables(self.rnn_model.variables)
self.rnn_model.summary()
@override(RecurrentNetwork)
rllib/examples/models/rnn_spy_model.py
-1
View File
@@ -107,7 +107,6 @@ class RNNSpyModel(RecurrentNetwork):
[inputs, seq_lens, state_in_h, state_in_c],
[logits, value_out, state_out_h, state_out_c])
self.base_model.summary()
self.register_variables(self.base_model.variables)
@override(RecurrentNetwork)
def forward_rnn(self, inputs, state, seq_lens):
rllib/examples/models/shared_weights_model.py
-3
View File
@@ -39,7 +39,6 @@ class TF2SharedWeightsModel(TFModelV2):
vf = tf.keras.layers.Dense(
units=1, activation=None, name="value_out")(last_layer)
self.base_model = tf.keras.models.Model(inputs, [output, vf])
self.register_variables(self.base_model.variables)
@override(ModelV2)
def forward(self, input_dict, state, seq_lens):
@@ -80,7 +79,6 @@ class SharedWeightsModel1(TFModelV2):
vf = tf.keras.layers.Dense(
units=1, activation=None, name="value_out")(last_layer)
self.base_model = tf.keras.models.Model(inputs, [output, vf])
self.register_variables(self.base_model.variables)
@override(ModelV2)
def forward(self, input_dict, state, seq_lens):
@@ -114,7 +112,6 @@ class SharedWeightsModel2(TFModelV2):
vf = tf.keras.layers.Dense(
units=1, activation=None, name="value_out")(last_layer)
self.base_model = tf.keras.models.Model(inputs, [output, vf])
self.register_variables(self.base_model.variables)
@override(ModelV2)
def forward(self, input_dict, state, seq_lens):
rllib/examples/models/simple_rpg_model.py
-1
View File
@@ -47,7 +47,6 @@ class CustomTFRPGModel(TFModelV2):
name)
self.model = TFFCNet(obs_space, action_space, num_outputs,
model_config, name)
self.register_variables(self.model.variables())
def forward(self, input_dict, state, seq_lens):
# The unpacked input tensors, where M=MAX_PLAYERS, N=MAX_ITEMS:
rllib/examples/models/trajectory_view_utilizing_models.py
-1
View File
@@ -36,7 +36,6 @@ class FrameStackingCartPoleModel(TFModelV2):
out = tf.keras.layers.Dense(self.num_outputs)(layer1)
values = tf.keras.layers.Dense(1)(layer1)
self.base_model = tf.keras.models.Model([input_], [out, values])
self.register_variables(self.base_model.variables)
self._last_value = None
rllib/models/catalog.py
+24 -14
View File
@@ -390,7 +390,9 @@ class ModelCatalog:
if model_config.get("use_lstm") else AttentionWrapper)
model_cls._wrapped_forward = forward
# Track and warn if vars were created but not registered.
# Obsolete: Track and warn if vars were created but not
# registered. Only still do this, if users do register their
# variables. If not (which they shouldn't), don't check here.
created = set()
def track_var_creation(next_creator, **kw):
@@ -419,19 +421,27 @@ class ModelCatalog:
# Other error -> re-raise.
else:
raise e
registered = set(instance.variables())
not_registered = set()
for var in created:
if var not in registered:
not_registered.add(var)
if not_registered:
raise ValueError(
"It looks like variables {} were created as part "
"of {} but does not appear in model.variables() "
"({}). Did you forget to call "
"model.register_variables() on the variables in "
"question?".format(not_registered, instance,
registered))
# User still registered TFModelV2's variables: Check, whether
# ok.
registered = set(instance.var_list)
if len(registered) > 0:
not_registered = set()
for var in created:
if var not in registered:
not_registered.add(var)
if not_registered:
raise ValueError(
"It looks like you are still using "
"`{}.register_variables()` to register your "
"model's weights. This is no longer required, but "
"if you are still calling this method at least "
"once, you must make sure to register all created "
"variables properly. The missing variables are {},"
" and you only registered {}. "
"Did you forget to call `register_variables()` on "
"some of the variables in question?".format(
instance, not_registered, registered))
elif framework == "torch":
# Try wrapping custom model with LSTM/attention, if required.
if model_config.get("use_lstm") or \
rllib/models/tests/test_distributions.py
+3 -6
View File
@@ -94,8 +94,7 @@ class TestDistributions(unittest.TestCase):
inputs = inputs_space.sample()
for fw, sess in framework_iterator(
session=True, frameworks=("tf", "tf2", "torch")):
for fw, sess in framework_iterator(session=True):
# Create the correct distribution object.
cls = JAXCategorical if fw == "jax" else Categorical if \
fw != "torch" else TorchCategorical
@@ -218,8 +217,7 @@ class TestDistributions(unittest.TestCase):
input_space = Box(-2.0, 2.0, shape=(2000, 10))
low, high = -2.0, 1.0
for fw, sess in framework_iterator(
frameworks=("torch", "tf", "tfe"), session=True):
for fw, sess in framework_iterator(session=True):
cls = SquashedGaussian if fw != "torch" else TorchSquashedGaussian
# Do a stability test using extreme NN outputs to see whether
@@ -310,8 +308,7 @@ class TestDistributions(unittest.TestCase):
"""Tests the DiagGaussian ActionDistribution for all frameworks."""
input_space = Box(-2.0, 1.0, shape=(2000, 10))
for fw, sess in framework_iterator(
frameworks=("torch", "tf", "tfe"), session=True):
for fw, sess in framework_iterator(session=True):
cls = DiagGaussian if fw != "torch" else TorchDiagGaussian
# Do a stability test using extreme NN outputs to see whether
rllib/models/tests/test_models.py
+59
View File
@@ -0,0 +1,59 @@
from gym.spaces import Box
import numpy as np
import unittest
from ray.rllib.models.tf.tf_modelv2 import TFModelV2
from ray.rllib.models.tf.fcnet import FullyConnectedNetwork
from ray.rllib.utils.framework import try_import_tf
tf1, tf, tfv = try_import_tf()
class TestTFModel(TFModelV2):
def __init__(self, obs_space, action_space, num_outputs, model_config,
name):
super().__init__(obs_space, action_space, num_outputs, model_config,
name)
input_ = tf.keras.layers.Input(shape=(3, ))
output = tf.keras.layers.Dense(2)(input_)
# A keras model inside.
self.keras_model = tf.keras.models.Model([input_], [output])
# A RLlib FullyConnectedNetwork (tf) inside (which is also a keras
# Model).
self.fc_net = FullyConnectedNetwork(obs_space, action_space, 3, {},
"fc1")
def forward(self, input_dict, state, seq_lens):
obs = input_dict["obs_flat"]
out1 = self.keras_model(obs)
out2, _ = self.fc_net({"obs": obs})
return tf.concat([out1, out2], axis=-1), []
class TestModels(unittest.TestCase):
"""Tests ModelV2 classes and their modularization capabilities."""
def test_tf_modelv2(self):
obs_space = Box(-1.0, 1.0, (3, ))
action_space = Box(-1.0, 1.0, (2, ))
my_tf_model = TestTFModel(obs_space, action_space, 5, {},
"my_tf_model")
# Call the model.
out, states = my_tf_model({"obs": np.array([obs_space.sample()])})
self.assertTrue(out.shape == (1, 5))
self.assertTrue(out.dtype == tf.float32)
self.assertTrue(states == [])
vars = my_tf_model.variables(as_dict=True)
self.assertTrue(len(vars) == 6)
self.assertTrue("keras_model.dense.kernel:0" in vars)
self.assertTrue("keras_model.dense.bias:0" in vars)
self.assertTrue("fc_net.base_model.fc_out.kernel:0" in vars)
self.assertTrue("fc_net.base_model.fc_out.bias:0" in vars)
self.assertTrue("fc_net.base_model.value_out.kernel:0" in vars)
self.assertTrue("fc_net.base_model.value_out.bias:0" in vars)
if __name__ == "__main__":
import pytest
import sys
sys.exit(pytest.main(["-v", __file__]))
rllib/models/tf/attention_net.py
-5
View File
@@ -117,7 +117,6 @@ class TrXLNet(RecurrentNetwork):
name="logits")(E_out)
self.base_model = tf.keras.models.Model([inputs], [logits])
self.register_variables(self.base_model.variables)
@override(RecurrentNetwork)
def forward_rnn(self, inputs: TensorType, state: List[TensorType],
@@ -287,7 +286,6 @@ class GTrXLNet(RecurrentNetwork):
self.trxl_model = tf.keras.Model(
inputs=[input_layer] + memory_ins, outputs=outs + memory_outs[:-1])
self.register_variables(self.trxl_model.variables)
self.trxl_model.summary()
# __sphinx_doc_begin__
@@ -386,7 +384,6 @@ class AttentionWrapper(TFModelV2):
position_wise_mlp_dim=cfg["attention_position_wise_mlp_dim"],
init_gru_gate_bias=cfg["attention_init_gru_gate_bias"],
)
self.register_variables(self.gtrxl.variables())
# `self.num_outputs` right now is the number of nodes coming from the
# attention net.
@@ -399,11 +396,9 @@ class AttentionWrapper(TFModelV2):
# values.
out = tf.keras.layers.Dense(self.num_outputs, activation=None)(input_)
self._logits_branch = tf.keras.models.Model([input_], [out])
self.register_variables(self._logits_branch.variables)
out = tf.keras.layers.Dense(1, activation=None)(input_)
self._value_branch = tf.keras.models.Model([input_], [out])
self.register_variables(self._value_branch.variables)
self.view_requirements = self.gtrxl.view_requirements
rllib/models/tf/fcnet.py
-2
View File
@@ -33,7 +33,6 @@ class FullyConnectedNetwork(TFModelV2):
num_outputs = num_outputs // 2
self.log_std_var = tf.Variable(
[0.0] * num_outputs, dtype=tf.float32, name="log_std")
self.register_variables([self.log_std_var])
# We are using obs_flat, so take the flattened shape as input.
inputs = tf.keras.layers.Input(
@@ -115,7 +114,6 @@ class FullyConnectedNetwork(TFModelV2):
self.base_model = tf.keras.Model(
inputs, [(logits_out
if logits_out is not None else last_layer), value_out])
self.register_variables(self.base_model.variables)
def forward(self, input_dict: Dict[str, TensorType],
state: List[TensorType],
rllib/models/tf/recurrent_net.py
-2
View File
@@ -50,7 +50,6 @@ class RecurrentNetwork(TFModelV2):
self.rnn_model = tf.keras.Model(
inputs=[input_layer, seq_in, state_in_h, state_in_c],
outputs=[output_layer, state_h, state_c])
self.register_variables(self.rnn_model.variables)
self.rnn_model.summary()
"""
@@ -179,7 +178,6 @@ class LSTMWrapper(RecurrentNetwork):
self._rnn_model = tf.keras.Model(
inputs=[input_layer, seq_in, state_in_h, state_in_c],
outputs=[logits, values, state_h, state_c])
self.register_variables(self._rnn_model.variables)
self._rnn_model.summary()
# Add prev-a/r to this model's view, if required.
rllib/models/tf/tf_modelv2.py
+67 -10
View File
@@ -1,9 +1,12 @@
import contextlib
import gym
from typing import List
import re
from typing import Dict, List, Union
from ray.util import log_once
from ray.rllib.models.modelv2 import ModelV2
from ray.rllib.utils.annotations import override, PublicAPI
from ray.rllib.utils.deprecation import deprecation_warning
from ray.rllib.utils.framework import try_import_tf
from ray.rllib.utils.typing import ModelConfigDict, TensorType
@@ -12,7 +15,7 @@ tf1, tf, tfv = try_import_tf()
@PublicAPI
class TFModelV2(ModelV2):
"""TF version of ModelV2.
"""TF version of ModelV2, which is always also a keras Model.
Note that this class by itself is not a valid model unless you
implement forward() in a subclass."""
@@ -33,18 +36,18 @@ class TFModelV2(ModelV2):
value_layer = tf.keras.layers.Dense(...)(hidden_layer)
self.base_model = tf.keras.Model(
input_layer, [output_layer, value_layer])
self.register_variables(self.base_model.variables)
"""
ModelV2.__init__(
self,
super().__init__(
obs_space,
action_space,
num_outputs,
model_config,
name,
framework="tf")
# Deprecated: TFModelV2 now automatically track their variables.
self.var_list = []
if tf1.executing_eagerly():
self.graph = None
else:
@@ -65,19 +68,73 @@ class TFModelV2(ModelV2):
def register_variables(self, variables: List[TensorType]) -> None:
"""Register the given list of variables with this model."""
if log_once("deprecated_tfmodelv2_register_variables"):
deprecation_warning(
old="TFModelV2.register_variables", error=False)
self.var_list.extend(variables)
@override(ModelV2)
def variables(self, as_dict: bool = False) -> List[TensorType]:
def variables(self, as_dict: bool = False) -> \
Union[List[TensorType], Dict[str, TensorType]]:
if as_dict:
return {v.name: v for v in self.var_list}
return list(self.var_list)
# Old way using `register_variables`.
if self.var_list:
return {v.name: v for v in self.var_list}
# New way: Automatically determine the var tree.
else:
return self._find_sub_modules("", self.__dict__)
# Old way using `register_variables`.
if self.var_list:
return list(self.var_list)
# New way: Automatically determine the var tree.
else:
return list(self.variables(as_dict=True).values())
@override(ModelV2)
def trainable_variables(self, as_dict: bool = False) -> List[TensorType]:
def trainable_variables(self, as_dict: bool = False) -> \
Union[List[TensorType], Dict[str, TensorType]]:
if as_dict:
return {
k: v
for k, v in self.variables(as_dict=True).items() if v.trainable
}
return [v for v in self.variables() if v.trainable]
@staticmethod
def _find_sub_modules(current_key, struct):
# Keras Model: key=k + "." + var-name (replace '/' by '.').
if isinstance(struct, tf.keras.models.Model):
ret = {}
for var in struct.variables:
key = current_key + "." + re.sub("/", ".", var.name)
ret[key] = var
return ret
# Other TFModelV2: Include its vars into ours.
elif isinstance(struct, TFModelV2):
return {
current_key + "." + key: var
for key, var in struct.variables(as_dict=True).items()
}
# tf.Variable
elif isinstance(struct, tf.Variable):
return {current_key + "." + struct.name: struct}
# List/Tuple.
elif isinstance(struct, (tuple, list)):
ret = {}
for i, value in enumerate(struct):
sub_vars = TFModelV2._find_sub_modules(
current_key + "_{}".format(i), value)
ret.update(sub_vars)
return ret
# Dict.
elif isinstance(struct, dict):
if current_key:
current_key += "_"
ret = {}
for key, value in struct.items():
sub_vars = TFModelV2._find_sub_modules(current_key + key,
value)
ret.update(sub_vars)
return ret
return {}
rllib/models/tf/visionnet.py
-1
View File
@@ -152,7 +152,6 @@ class VisionNetwork(TFModelV2):
lambda x: tf.squeeze(x, axis=[1, 2]))(last_layer)
self.base_model = tf.keras.Model(inputs, [conv_out, value_out])
self.register_variables(self.base_model.variables)
# Optional: framestacking obs/new_obs for Atari.
if self.traj_view_framestacking:
rllib/models/torch/torch_modelv2.py
+5 -3
View File
@@ -1,5 +1,5 @@
import gym
from typing import List
from typing import Dict, List, Union
from ray.rllib.models.modelv2 import ModelV2
from ray.rllib.utils.annotations import override, PublicAPI
@@ -47,14 +47,16 @@ class TorchModelV2(ModelV2):
framework="torch")
@override(ModelV2)
def variables(self, as_dict: bool = False) -> List[TensorType]:
def variables(self, as_dict: bool = False) -> \
Union[List[TensorType], Dict[str, TensorType]]:
p = list(self.parameters())
if as_dict:
return {k: p[i] for i, k in enumerate(self.state_dict().keys())}
return p
@override(ModelV2)
def trainable_variables(self, as_dict: bool = False) -> List[TensorType]:
def trainable_variables(self, as_dict: bool = False) -> \
Union[List[TensorType], Dict[str, TensorType]]:
if as_dict:
return {
k: v
rllib/tests/test_model_imports.py
-2
View File
@@ -49,8 +49,6 @@ class MyKerasModel(TFModelV2):
else:
self.base_model = tf.keras.Model(self.inputs, layer_out)
self.register_variables(self.base_model.variables)
def forward(self, input_dict, state, seq_lens):
if self.model_config["vf_share_layers"]:
model_out, self._value_out = self.base_model(input_dict["obs"])
rllib/tests/test_nested_observation_spaces.py
-2
View File
@@ -240,7 +240,6 @@ class DictSpyModel(TFModelV2):
self.num_outputs = num_outputs or 64
out = tf.keras.layers.Dense(self.num_outputs)(input_)
self._main_layer = tf.keras.models.Model([input_], [out])
self.register_variables(self._main_layer.variables)
def forward(self, input_dict, state, seq_lens):
def spy(pos, front_cam, task):
@@ -282,7 +281,6 @@ class TupleSpyModel(TFModelV2):
self.num_outputs = num_outputs or 64
out = tf.keras.layers.Dense(self.num_outputs)(input_)
self._main_layer = tf.keras.models.Model([input_], [out])
self.register_variables(self._main_layer.variables)
def forward(self, input_dict, state, seq_lens):
def spy(pos, cam, task):
rllib/utils/deprecation.py
+9 -5
View File
@@ -1,4 +1,5 @@
import logging
from typing import Optional, Union
logger = logging.getLogger(__name__)
@@ -8,15 +9,18 @@ logger = logging.getLogger(__name__)
DEPRECATED_VALUE = -1
def deprecation_warning(old, new=None, error=None):
"""
Logs (via the `logger` object) or throws a deprecation warning/error.
def deprecation_warning(
old: str,
new: Optional[str] = None,
error: Optional[Union[bool, Exception]] = None) -> None:
"""Warns (via the `logger` object) or throws a deprecation warning/error.
Args:
old (str): A description of the "thing" that is to be deprecated.
new (Optional[str]): A description of the new "thing" that replaces it.
error (Optional[Union[bool,Exception]]): Whether or which exception to
throw. If True, throw ValueError.
error (Optional[Union[bool, Exception]]): Whether or which exception to
throw. If True, throw ValueError. If False, just warn.
If Exception, throw that Exception.
"""
msg = "`{}` has been deprecated.{}".format(
old, (" Use `{}` instead.".format(new) if new else ""))
rllib/utils/exploration/curiosity.py
-1
View File
@@ -208,7 +208,6 @@ class Curiosity(Exploration):
self._curiosity_feature_net.base_model.variables + \
self._curiosity_inverse_fcnet.variables + \
self._curiosity_forward_fcnet.variables
self.model.register_variables(self._optimizer_var_list)
self._optimizer = tf1.train.AdamOptimizer(learning_rate=self.lr)
# Create placeholders and initialize the loss.
if self.framework == "tf":