From a5d1f0351577a26395020dce9f8c621d185147ad Mon Sep 17 00:00:00 2001 From: Eugene Vinitsky Date: Sun, 13 Jan 2019 16:54:23 -0700 Subject: [PATCH] [rllib] fix for rollout of lstm policies (#3643) * fix for lstm policies * added call to local evaluator * Update python/ray/rllib/rollout.py Co-Authored-By: eugenevinitsky * Update rollout.py * Update rollout.py --- python/ray/rllib/rollout.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/python/ray/rllib/rollout.py b/python/ray/rllib/rollout.py index 09960f2bb..05f19e1e8 100755 --- a/python/ray/rllib/rollout.py +++ b/python/ray/rllib/rollout.py @@ -104,6 +104,17 @@ def rollout(agent, env_name, num_steps, out=None, no_render=True): env = agent.local_evaluator.env else: env = gym.make(env_name) + + if hasattr(agent, "local_evaluator"): + state_init = agent.local_evaluator.policy_map[ + "default"].get_initial_state() + else: + state_init = [] + if state_init: + use_lstm = True + else: + use_lstm = False + if out is not None: rollouts = [] steps = 0 @@ -114,7 +125,11 @@ def rollout(agent, env_name, num_steps, out=None, no_render=True): done = False reward_total = 0.0 while not done and steps < (num_steps or steps + 1): - action = agent.compute_action(state) + if use_lstm: + action, state_init, logits = agent.compute_action( + state, state=state_init) + else: + action = agent.compute_action(state) next_state, reward, done, _ = env.step(action) reward_total += reward if not no_render: