action="store_true", help="Init Ray in local mode for easier debugging.", ) args = parser.parse_args() print(f"Running with following CLI args: {args}") return args if __name__ == "__main__": args = get_cli_args() ray.init(num_cpus=args.num_cpus or None, local_mode=args.local_mode) # register custom environments registry.register_env("RepeatAfterMeEnv", lambda c: RepeatAfterMeEnv(c)) registry.register_env("RepeatInitialObsEnv", lambda _: RepeatInitialObsEnv()) registry.register_env("LookAndPush", lambda _: OneHot(LookAndPush())) registry.register_env("StatelessCartPole", lambda _: StatelessCartPole()) # main part: RLlib config with AttentionNet model config = { "env": args.env, # This env_config is only used for the RepeatAfterMeEnv env. "env_config": { "repeat_delay": 2, }, "gamma": 0.99, # Use GPUs iff `RLLIB_NUM_GPUS` env var set to > 0. "num_gpus": int(os.environ.get("RLLIB_NUM_GPUS", 0)),
parser.add_argument("--env", type=str, default="RepeatAfterMeEnv") parser.add_argument("--num-cpus", type=int, default=0) parser.add_argument("--as-test", action="store_true") parser.add_argument("--torch", action="store_true") parser.add_argument("--stop-reward", type=float, default=90) parser.add_argument("--stop-iters", type=int, default=100) parser.add_argument("--stop-timesteps", type=int, default=100000) if __name__ == "__main__": args = parser.parse_args() ray.init(num_cpus=args.num_cpus or None) ModelCatalog.register_custom_model( "rnn", TorchRNNModel if args.torch else RNNModel) register_env("RepeatAfterMeEnv", lambda c: RepeatAfterMeEnv(c)) register_env("RepeatInitialObsEnv", lambda _: RepeatInitialObsEnv()) config = { "env": args.env, "env_config": { "repeat_delay": 2, }, "gamma": 0.9, # Use GPUs iff `RLLIB_NUM_GPUS` env var set to > 0. "num_gpus": int(os.environ.get("RLLIB_NUM_GPUS", "0")), "num_workers": 0, "num_envs_per_worker": 20, "entropy_coeff": 0.001, "num_sgd_iter": 5, "vf_loss_coeff": 1e-5,
self._cur_value = torch.reshape(self.value_branch(lstm_out[0]), [-1]) return action_out, [ torch.squeeze(lstm_out[1][0], 0), torch.squeeze(lstm_out[1][1], 0) ] if __name__ == "__main__": args = parser.parse_args() ray.init(num_cpus=args.num_cpus or None) ModelCatalog.register_custom_model("rnn", RNNModel) tune.register_env("repeat_initial", lambda _: RepeatInitialObsEnv(episode_len=100)) tune.register_env("repeat_after_me", lambda _: RepeatAfterMeEnv({"repeat_delay": 1})) tune.register_env("stateless_cartpole", lambda _: StatelessCartPole()) config = { "env": args.env, "use_pytorch": True, "num_workers": 0, "num_envs_per_worker": 20, "gamma": 0.9, "entropy_coeff": 0.0001, "model": { "custom_model": "rnn", "max_seq_len": 20, "lstm_use_prev_action_reward": "store_true", "custom_options": { "fc_size": args.fc_size,