Spaces:
Sleeping
Sleeping
File size: 5,200 Bytes
403f471 | 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 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | import numpy as np
import torch
from itertools import chain
class Evaluator(object):
def __init__(self, environment, model, logger, agents, max_steps):
self.env = environment
self.model = model
self.logger = logger
self.agents = agents
self.max_steps = max_steps
def play_n_episodes(self, render=False, fixed_spawn=None, silent=False):
"""
wraps play_one_episode, playing a single episode at a time and logs
results used when playing demos.
"""
if fixed_spawn is None:
num_runs = 1
fixed_spawn = [None]
else:
# fixed_spawn should be, for example, [0.5 , 0.5 , 0.5, 0, 0, 0] for 2 runs
# In the first run agents spawn in the middle and in the second they will spawn from the corner
fixed_spawn = np.array(fixed_spawn).reshape((-1, 3)) # 3 dimensions
num_runs = fixed_spawn.shape[0]
# Set all the agents to the same spawn point
fixed_spawn = np.stack([fixed_spawn for _ in range(self.agents)], axis=-1)
num_files = self.env.files.num_files
self.model.train(False)
headers = ["number"] + list(chain.from_iterable(zip(
[f"Filename {i}" for i in range(self.agents)],
[f"Agent {i} pos x" for i in range(self.agents)],
[f"Agent {i} pos y" for i in range(self.agents)],
[f"Agent {i} pos z" for i in range(self.agents)],
[f"Landmark {i} pos x" for i in range(self.agents)],
[f"Landmark {i} pos y" for i in range(self.agents)],
[f"Landmark {i} pos z" for i in range(self.agents)],
[f"Distance {i}" for i in range(self.agents)])))
self.logger.write_locations(headers)
distances = []
for j in range(num_runs):
for k in range(num_files):
score, start_dists, q_values, info = self.play_one_episode(render, fixed_spawn=fixed_spawn[j])
row = [j * num_files + k + 1] + list(chain.from_iterable(zip(
[info[f"filename_{i}"] for i in range(self.agents)],
[info[f"agent_xpos_{i}"] for i in range(self.agents)],
[info[f"agent_ypos_{i}"] for i in range(self.agents)],
[info[f"agent_zpos_{i}"] for i in range(self.agents)],
[info.get(f"landmark_xpos_{i}", "N/A") for i in range(self.agents)],
[info.get(f"landmark_ypos_{i}", "N/A") for i in range(self.agents)],
[info.get(f"landmark_zpos_{i}", "N/A") for i in range(self.agents)],
[info.get(f"distError_{i}", "N/A") for i in range(self.agents)])))
for i in range(self.agents):
key = f"distError_{i}"
if key in info:
distances.append(info[key])
self.logger.write_locations(row)
if len(distances) == 0:
return None, None # No distance mean and std for task "play" as there is no ground truth
mean = np.mean(distances, 0)
std = np.std(distances, 0, ddof=1)
if not silent:
self.logger.log(f"mean distances {mean}")
self.logger.log(f"Std distances {std}")
return mean, std
def play_one_episode(self, render=False, frame_history=4, fixed_spawn=None):
device = next(self.model.parameters()).device
def predict(obs_stack):
"""
Run a full episode, mapping observation to action,
using greedy policy.
"""
inputs = torch.from_numpy(obs_stack).float().permute(
0, 4, 1, 2, 3).unsqueeze(0).to(device)
with torch.no_grad():
q_vals = self.model(inputs)
idx = torch.max(q_vals, -1)[1]
greedy_steps = np.array(idx, dtype=np.int32).flatten()
return greedy_steps, q_vals.detach().cpu().numpy().squeeze(0)
obs_stack = self.env.reset(fixed_spawn)
# Here obs have shape (agent, *image_size, frame_history)
sum_r = np.zeros((self.agents))
isOver = [False] * self.agents
start_dists = None
steps = 0
while steps < self.max_steps and not np.all(isOver):
acts, q_values = predict(obs_stack)
obs_stack, r, isOver, info = self.env.step(acts, q_values, isOver)
steps += 1
# Machine-parseable trajectory log line (agent 0 only). Consumed
# by app.py to render the agent's search path as an animation.
print("STEP_LOC: {} {} {} {}".format(
steps,
info.get('agent_xpos_0', 'NA'),
info.get('agent_ypos_0', 'NA'),
info.get('agent_zpos_0', 'NA'),
))
if start_dists is None:
start_dists = [
info.get('distError_' + str(i), "N/A") for i in range(self.agents)]
if render:
self.env.render()
for i in range(self.agents):
if not isOver[i]:
sum_r[i] += r[i]
return sum_r, start_dists, q_values, info
|