File size: 2,285 Bytes
529590a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 | def eval_one_episode(TASK_ENV, model_client):
model_client.call(func_name="reset") # reset policy
while not TASK_ENV.is_episode_end(): # Check whether the episode ends
obs = TASK_ENV.get_obs() # Get Observation
model_client.call(func_name="update_obs", obs=obs) # Update Observation
actions = model_client.call(func_name="get_action") # Get Action according to observation chunk
for action_idx, action in enumerate(actions):
TASK_ENV.take_action(action)
if TASK_ENV.is_episode_end() or action_idx + 1 == len(actions):
break
obs = TASK_ENV.get_obs()
model_client.call(func_name="update_obs", obs=obs)
def eval_one_episode_batch(TASK_ENV, model_client):
model_client.call(func_name="reset")
while not TASK_ENV.is_episode_end(): # Check whether the episode ends
env_idx_list = TASK_ENV.get_running_env_idx_list() # Get Running Environment Index List
obs_list = TASK_ENV.get_obs_batch(env_idx_list) # Get Observation
model_client.call(func_name="update_obs_batch", obs=obs_list)
actions = model_client.call(func_name="get_action_batch", obs=env_idx_list) # Get Action according to observation chunk
chunk_size = len(actions[0]) # Get the chunk size
for action_idx in range(chunk_size): # Iterate over the action chunk
current_action_list = [env_actions[action_idx] for env_actions in actions] # Get the current action list
TASK_ENV.take_action_batch(current_action_list, env_idx_list) # Take the action
if TASK_ENV.is_episode_end() or action_idx + 1 == chunk_size: # Check whether the episode ends
break
running = set(TASK_ENV.get_running_env_idx_list()) # Get the running environment index list
active_batch_idx = [i for i, env_idx in enumerate(env_idx_list) if env_idx in running] # Get the active batch index
actions = [actions[i] for i in active_batch_idx] # Get the active action list
env_idx_list = [env_idx_list[i] for i in active_batch_idx] # Get the active environment index list
model_client.call(func_name="update_obs_batch", obs=TASK_ENV.get_obs_batch(env_idx_list)) # Update the observation |