def evaluate(logger, log_path):
retina = Retina()
lip = LIP()
vc = VC()
pfc = PFC()
fef = FEF()
bg = BG()
sc = SC()
hp = HP()
cb = CB()
agent = Agent(retina=retina,
lip=lip,
vc=vc,
pfc=pfc,
fef=fef,
bg=bg,
sc=sc,
hp=hp,
cb=cb)
pfc.load_model('data/pfc_task_detection.pth')
#bg.load_model('data/bg_rl.pth')
total_reward = 0
all_trial_results = []
for i, task in enumerate(tasks):
trial_results, task_reward = task.evaluate(agent)
all_trial_results += trial_results
total_reward += task_reward
logger.log("evaluation_reward", total_reward, i)
# Save result csv files
save_results(all_trial_results, log_path)
print("evaluation finished:")
logger.close()
class Inspector(object):
def __init__(self, content, display_size):
pygame.init()
self.surface = pygame.display.set_mode(display_size, 0, 24)
pygame.display.set_caption('oculomotor')
self.retina = Retina()
self.lip = LIP()
self.vc = VC()
self.pfc = PFC()
self.fef = FEF()
self.bg = BG()
self.sc = SC()
self.hp = HP()
self.cb = CB()
self.agent = Agent(
retina=self.retina,
lip=self.lip,
vc=self.vc,
pfc=self.pfc,
fef=self.fef,
bg=self.bg,
sc=self.sc,
hp=self.hp,
cb=self.cb
)
self.env = Environment(content)
self.pfc.load_model('data/pfc_task_detection.pth')
#self.bg.load_model('data/bg_rl.pth')
obs = self.env.reset()
self.last_image = obs['screen']
self.last_angle = obs['angle']
self.last_reward = 0
self.last_done = False
self.episode_reward = 0
self.font = pygame.font.Font(None, 20)
self.display_size = display_size
def update(self):
self.surface.fill(BLACK)
done = self.process()
pygame.display.update()
return done
def draw_text(self, str, left, top, color=WHITE):
text = self.font.render(str, True, color, BLACK)
text_rect = text.get_rect()
text_rect.left = left
text_rect.top = top
self.surface.blit(text, text_rect)
def draw_center_text(self, str, center_x, top):
text = self.font.render(str, True, WHITE, BLACK)
text_rect = text.get_rect()
text_rect.centerx = center_x
text_rect.top = top
self.surface.blit(text, text_rect)
def show_rgb_256_image(self, data, left, top, label):
self.show_image((data * 1.0).astype(np.uint8), left, top, label)
def show_gray_1_image(self, data, left, top, label):
data = np.clip(data * 255.0, 0.0, 255.0)
data = data.astype(np.uint8)
data = np.stack([data for _ in range(3)], axis=2)
self.show_image(data, left, top, label)
def show_optical_flow(self, optical_flow, left, top, label):
# Show optical flow with HSV color image
image = self.get_optical_flow_hsv(optical_flow)
# Draw optical flow direction with lines
step = 16
h, w = optical_flow.shape[:2]
y, x = np.mgrid[step // 2:h:step, step // 2:w:step].reshape(
2, -1).astype(int)
fx, fy = optical_flow[y, x].T * 10
lines = np.vstack([x, y, x + fx, y + fy]).T.reshape(-1, 2, 2)
lines = np.int32(lines + 0.5)
cv2.polylines(image, lines, 0, (0, 255, 0))
for (x1, y1), (x2, y2) in lines:
cv2.circle(image, (x1, y1), 1, (0, 255, 0), -1)
self.show_image(image, left, top, label)
def get_optical_flow_hsv(self, optical_flow):
h, w = optical_flow.shape[:2]
fx, fy = optical_flow[:, :, 0], optical_flow[:, :, 1]
ang = np.arctan2(fy, fx) + np.pi
v = np.sqrt(fx * fx + fy * fy)
hsv = np.zeros((h, w, 3), np.uint8)
hsv[..., 0] = ang * (180 / np.pi / 2)
hsv[..., 1] = 255
hsv[..., 2] = np.minimum(v * 4, 255)
image = cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB)
return image
def show_image(self, data, left, top, label):
image = pygame.image.frombuffer(data, (128, 128), 'RGB')
self.surface.blit(image, (left, top))
self.draw_center_text(label, 128 / 2 + left, top + 128 + 8)
pygame.draw.rect(self.surface, DARK_GRAY, Rect(left, top, 128, 128), 1)
def show_reward(self):
self.draw_text("TASK: {}, REWARD: {}, PHASE: {}, PFC_STEPS: {}, ENV_STEPS: {}".format(self.pfc.task, int(self.episode_reward), self.pfc.phase, self.pfc.steps, self.env.content.step_count),
8, 434 + 48)
def show_fef_data_bars(self, fef_data):
fef_data_len = len(fef_data)
bottom = 256 + 16
pygame.draw.line(self.surface, DARK_GRAY, (8, bottom - 100),
(3 * fef_data_len + 8, bottom - 100), 1)
for i, data in enumerate(fef_data):
likelihood = data[0]
left = 8 + 3 * i
top = bottom - 100 * likelihood
pygame.draw.line(self.surface, WHITE, (left, top), (left, bottom),
1)
self.draw_center_text("likelihoods", (8 + 3 * fef_data_len) // 4,
bottom + 8)
def show_fef_data_grid(self, fef_data):
grid_division = int(math.sqrt(len(fef_data) // 4))
grid_width = 128 // grid_division
likelihoods0 = []
likelihoods1 = []
likelihoods2 = []
likelihoods3 = []
data_len = len(fef_data) // 4
for i in range(data_len):
likelihoods0.append(fef_data[i][0])
likelihoods1.append(fef_data[i + data_len][0])
likelihoods2.append(fef_data[i + data_len*2][0])
likelihoods3.append(fef_data[i + data_len*3][0])
if self.pfc.task in [6]:
self.show_grid(likelihoods1, 0, grid_division, grid_width, 8 + 128*0, 330, "direction acc")
if self.pfc.task in [1, 4, 5]:
self.show_grid(likelihoods0, 0, grid_division, grid_width, 8 + 128*1, 330, "template acc")
if self.pfc.task in [2]:
self.show_grid(likelihoods2, 0, grid_division, grid_width, 8 + 128*2, 330, "change acc")
if self.pfc.task in [1, 3, 4, 5]:
self.show_grid(likelihoods3, 0, grid_division, grid_width, 8 + 128*3, 330, "search acc")
def show_grid(self, data, offset, grid_division, grid_width, left, top,
label):
index = 0
for ix in range(grid_division):
x = grid_width * ix
for iy in range(grid_division):
y = grid_width * iy
likelihood = data[index]
c = int(likelihood * 255.0)
color = (c, c, c)
pygame.draw.rect(self.surface, color,
Rect(left + x, top + y, grid_width,
grid_width))
index += 1
pygame.draw.rect(self.surface, DARK_GRAY, Rect(left, top, 128, 128), 1)
self.draw_center_text(label, 128 / 2 + left, top + 128 + 8)
def process(self):
action = self.agent(self.last_image, self.last_angle, self.last_reward,
self.last_done)
obs, reward, done, _ = self.env.step(action)
self.episode_reward += reward
if done:
obs = self.env.reset()
self.episode_reward = 0
image = obs['screen']
angle = obs['angle']
self.show_rgb_256_image(image, 128*0+8, 128*0+8, 'input')
if self.retina.last_retina_image is not None:
self.show_rgb_256_image(self.retina.last_retina_image, 128*1+8, 128*0+8, 'retina')
if self.pfc.last_memory_image is not None:
self.show_rgb_256_image(self.pfc.last_memory_image, 128*2+8, 128*0+8, 'working memory')
if self.pfc.last_saliency_map is not None:
self.show_gray_1_image(self.pfc.last_saliency_map, 128*3 + 8, 128*0+8, 'saliency map')
if self.pfc.task in [6] and self.lip.last_optical_flow is not None:
self.show_optical_flow(self.lip.last_optical_flow, 128*0+8, 128*1+8+32, "optical flow")
if self.pfc.task in [2] and self.pfc.last_change_map is not None:
self.show_rgb_256_image(self.pfc.last_change_map, 128*2+8, 128*1+8+32, 'change map')
if self.pfc.task in [1, 3, 4, 5] and self.pfc.last_search_map is not None:
self.show_gray_1_image(self.pfc.last_search_map, 128*3+8, 128*1+8+32, 'search map')
#if self.lip.last_small_e_match_map is not None:
# self.show_gray_1_image(self.lip.last_small_e_match_map, 128 * 3 + 8, 8, 'match small E')
#if self.lip.last_magenta_t_match_map is not None:
# self.show_gray_1_image(self.lip.last_magenta_t_match_map, 128 * 3 + 8, 8, 'match magenta T')
if self.sc.last_fef_data is not None:
#self.show_fef_data_bars(self.sc.last_fef_data)
self.show_fef_data_grid(self.sc.last_fef_data)
#if self.hp.map_image is not None:
# self.show_rgb_256_image(self.hp.map_image, 128 * 3 + 8, 300, 'allocentric map')
self.show_reward()
self.last_image = image
self.last_angle = angle
self.last_reward = reward
self.last_done = done
return done
def get_frame(self):
data = self.surface.get_buffer().raw
image = np.fromstring(data, dtype=np.uint8)
image = image.reshape((self.display_size[1], self.display_size[0], 3))
return image
def train(content, step_size, logger):
retina = Retina()
lip = LIP()
vc = VC()
pfc = PFC()
fef = FEF()
bg = BG()
sc = SC()
hp = HP()
cb = CB()
agent = Agent(retina=retina,
lip=lip,
vc=vc,
pfc=pfc,
fef=fef,
bg=bg,
sc=sc,
hp=hp,
cb=cb)
env = Environment(content)
# If your training code is written inside BG module, please add model load code here like.
#
pfc.load_model('data/pfc_task_detection.pth')
#bg.load_model('data/bg_rl.pth')
#
# When runnning with Docker, directory under 'oculomotor/' is volume shared
# with the host, so you can load/save the model data at anywhere under 'oculomotor/' dir.
obs = env.reset()
reward = 0
done = False
episode_reward = 0
episode_count = 0
# Add initial reward log
logger.log("episode_reward", episode_reward, episode_count)
for i in range(step_size):
image, angle = obs['screen'], obs['angle']
# Choose action by the agent's decision
action = agent(image, angle, reward, done)
# Foward environment one step
obs, reward, done, _ = env.step(action)
episode_reward += reward
if done:
obs = env.reset()
print("episode reward={}".format(episode_reward))
# Store log for tensorboard graph
episode_count += 1
logger.log("episode_reward", episode_reward, episode_count)
episode_reward = 0
# Plase add model save code as you like.
#
# if i % 10 == 0:
#bg.save_model('data/bg_rl.pth')
print("training finished")
logger.close()