mirror of
https://github.com/vale981/ray
synced 2025-03-09 04:46:38 -04:00
60d4d5e1aa
* Remove all __future__ imports from RLlib. * Remove (object) again from tf_run_builder.py::TFRunBuilder. * Fix 2xLINT warnings. * Fix broken appo_policy import (must be appo_tf_policy) * Remove future imports from all other ray files (not just RLlib). * Remove future imports from all other ray files (not just RLlib). * Remove future import blocks that contain `unicode_literals` as well. Revert appo_tf_policy.py to appo_policy.py (belongs to another PR). * Add two empty lines before Schedule class. * Put back __future__ imports into determine_tests_to_run.py. Fails otherwise on a py2/print related error.
57 lines
1.7 KiB
Python
57 lines
1.7 KiB
Python
"""Example of using a custom model with batch norm."""
|
|
|
|
import argparse
|
|
|
|
import ray
|
|
from ray import tune
|
|
from ray.rllib.models import Model, ModelCatalog
|
|
from ray.rllib.models.tf.misc import normc_initializer
|
|
from ray.rllib.utils import try_import_tf
|
|
|
|
tf = try_import_tf()
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--num-iters", type=int, default=200)
|
|
parser.add_argument("--run", type=str, default="PPO")
|
|
|
|
|
|
class BatchNormModel(Model):
|
|
def _build_layers_v2(self, input_dict, num_outputs, options):
|
|
last_layer = input_dict["obs"]
|
|
hiddens = [256, 256]
|
|
for i, size in enumerate(hiddens):
|
|
label = "fc{}".format(i)
|
|
last_layer = tf.layers.dense(
|
|
last_layer,
|
|
size,
|
|
kernel_initializer=normc_initializer(1.0),
|
|
activation=tf.nn.tanh,
|
|
name=label)
|
|
# Add a batch norm layer
|
|
last_layer = tf.layers.batch_normalization(
|
|
last_layer, training=input_dict["is_training"])
|
|
output = tf.layers.dense(
|
|
last_layer,
|
|
num_outputs,
|
|
kernel_initializer=normc_initializer(0.01),
|
|
activation=None,
|
|
name="fc_out")
|
|
return output, last_layer
|
|
|
|
|
|
if __name__ == "__main__":
|
|
args = parser.parse_args()
|
|
ray.init()
|
|
|
|
ModelCatalog.register_custom_model("bn_model", BatchNormModel)
|
|
tune.run(
|
|
args.run,
|
|
stop={"training_iteration": args.num_iters},
|
|
config={
|
|
"env": "Pendulum-v0" if args.run == "DDPG" else "CartPole-v0",
|
|
"model": {
|
|
"custom_model": "bn_model",
|
|
},
|
|
"num_workers": 0,
|
|
},
|
|
)
|