ray/rllib/examples/cartpole_lstm.py
Sven Mika 42991d723f
[RLlib] rllib/examples folder restructuring (#8250)
Cleans up of the rllib/examples folder by moving all example Envs into rllibexamples/env (so they can be used by other scripts and tests as well).
2020-05-01 22:59:34 +02:00

45 lines
1.2 KiB
Python

import argparse
from ray.rllib.examples.env.stateless_cartpole import StatelessCartPole
parser = argparse.ArgumentParser()
parser.add_argument("--stop", type=int, default=200)
parser.add_argument("--torch", action="store_true")
parser.add_argument("--use-prev-action-reward", action="store_true")
parser.add_argument("--run", type=str, default="PPO")
parser.add_argument("--num-cpus", type=int, default=0)
if __name__ == "__main__":
import ray
from ray import tune
args = parser.parse_args()
ray.init(num_cpus=args.num_cpus or None)
configs = {
"PPO": {
"num_sgd_iter": 5,
"vf_share_layers": True,
"vf_loss_coeff": 0.0001,
},
"IMPALA": {
"num_workers": 2,
"num_gpus": 0,
"vf_loss_coeff": 0.01,
},
}
tune.run(
args.run,
stop={"episode_reward_mean": args.stop},
config=dict(
configs[args.run], **{
"env": StatelessCartPole,
"model": {
"use_lstm": True,
"lstm_use_prev_action_reward": args.use_prev_action_reward,
},
"use_pytorch": args.torch,
}),
)