obs_dict = {self.i: self.last_obs[self.i]}
self.i = (self.i + 1) % self.num
return obs_dict
def step(self, action_dict):
assert len(self.dones) != len(self.agents)
for i, action in action_dict.items():
(
self.last_obs[i],
self.last_rew[i],
self.last_done[i],
self.last_info[i],
) = self.agents[i].step(action)
obs = {self.i: self.last_obs[self.i]}
rew = {self.i: self.last_rew[self.i]}
done = {self.i: self.last_done[self.i]}
info = {self.i: self.last_info[self.i]}
if done[self.i]:
rew[self.i] = 0
self.dones.add(self.i)
self.i = (self.i + 1) % self.num
done["__all__"] = len(self.dones) == len(self.agents)
return obs, rew, done, info
MultiAgentCartPole = make_multi_agent("CartPole-v0")
MultiAgentMountainCar = make_multi_agent("MountainCarContinuous-v0")
MultiAgentPendulum = make_multi_agent("Pendulum-v1")
MultiAgentStatelessCartPole = make_multi_agent(
lambda config: StatelessCartPole(config))
parser.add_argument("--framework", choices=["tf", "torch"], default="tf")
parser.add_argument("--as-test", action="store_true")
parser.add_argument("--stop-iters", type=int, default=200)
parser.add_argument("--stop-timesteps", type=int, default=500000)
parser.add_argument("--stop-reward", type=float, default=80)
if __name__ == "__main__":
args = parser.parse_args()
ray.init(num_cpus=args.num_cpus or None)
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())
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)),
"num_workers": 0,
"num_envs_per_worker": 20,
"entropy_coeff": 0.001,
"num_sgd_iter": 10,
"vf_loss_coeff": 1e-5,
parser.add_argument("--framework",
choices=["tf2", "tf", "tfe", "torch"],
default="tf")
parser.add_argument("--as-test", action="store_true")
parser.add_argument("--stop-iters", type=int, default=50)
parser.add_argument("--stop-timesteps", type=int, default=200000)
parser.add_argument("--stop-reward", type=float, default=150.0)
if __name__ == "__main__":
args = parser.parse_args()
ray.init(num_cpus=3)
ModelCatalog.register_custom_model(
"frame_stack_model", FrameStackingCartPoleModel
if args.framework != "torch" else TorchFrameStackingCartPoleModel)
tune.register_env("stateless_cartpole", lambda c: StatelessCartPole())
config = {
"env": "stateless_cartpole",
"model": {
"vf_share_layers": True,
"custom_model": "frame_stack_model",
"custom_model_config": {
"num_frames": 16,
},
# To compare against a simple LSTM:
# "use_lstm": True,
# "lstm_use_prev_action": True,
# "lstm_use_prev_reward": True,
results = tune.run(
args.run, config=config, stop=stop, verbose=2, checkpoint_at_end=True)
if args.as_test:
check_learning_achieved(results, args.stop_reward)
checkpoints = results.get_trial_checkpoints_paths(
trial=results.get_best_trial("episode_reward_mean", mode="max"),
metric="episode_reward_mean")
checkpoint_path = checkpoints[0][0]
trainer = PPOTrainer(config)
trainer.restore(checkpoint_path)
# Inference loop.
env = StatelessCartPole()
# Run manual inference loop for n episodes.
for _ in range(10):
episode_reward = 0.0
reward = 0.0
action = 0
done = False
obs = env.reset()
while not done:
# Create a dummy action using the same observation n times,
# as well as dummy prev-n-actions and prev-n-rewards.
action, state, logits = trainer.compute_single_action(
input_dict={
"obs": obs,
"prev_n_obs": np.stack([obs for _ in range(num_frames)]),
self.last_info[i] = {}
obs_dict = {self.i: self.last_obs[self.i]}
self.i = (self.i + 1) % self.num
return obs_dict
def step(self, action_dict):
assert len(self.dones) != len(self.agents)
for i, action in action_dict.items():
(
self.last_obs[i],
self.last_rew[i],
self.last_done[i],
self.last_info[i],
) = self.agents[i].step(action)
obs = {self.i: self.last_obs[self.i]}
rew = {self.i: self.last_rew[self.i]}
done = {self.i: self.last_done[self.i]}
info = {self.i: self.last_info[self.i]}
if done[self.i]:
rew[self.i] = 0
self.dones.add(self.i)
self.i = (self.i + 1) % self.num
done["__all__"] = len(self.dones) == len(self.agents)
return obs, rew, done, info
MultiAgentCartPole = make_multi_agent("CartPole-v0")
MultiAgentMountainCar = make_multi_agent("MountainCarContinuous-v0")
MultiAgentPendulum = make_multi_agent("Pendulum-v1")
MultiAgentStatelessCartPole = make_multi_agent(lambda config: StatelessCartPole(config))