def run_episode(env, thing_below, goal_thing_below, nvm, init_regs, init_conns, penalty_tracker, sigma=0): # reload blocks env.reset() env.load_blocks(thing_below) # invert goals for nvm goal_thing_above = invert(goal_thing_below, num_blocks=len(thing_below), num_bases=len(env.bases)) for key, val in goal_thing_above.items(): if val == "none": goal_thing_above[key] = "nil" # reset nvm, input new env, mount main program nvm.reset_state(init_regs, init_conns) memorize_env(nvm, goal_thing_above) nvm.mount("main") log_prob = 0.0 # accumulate over episode log_probs, rewards = [], [] dbg = False if dbg: nvm.dbg() target_changed = False while True: done = nvm.tick() # reliable if core is not trained if dbg: nvm.dbg() # if nvm.tick_counter % 100 == 0: print(" tick %d" % nvm.tick_counter) if target_changed: mu = nvm.registers["jnt"].content if sigma > 0: dist = tr.distributions.normal.Normal(mu, sigma) position = dist.sample() log_probs.append( dist.log_prob(position).sum()) # multivariate white noise log_prob += log_probs[-1] else: position = mu penalty_tracker.reset() # nvm.dbg() # print(" pos:", position.detach().numpy()) nvm.env.goto_position(position.detach().numpy()) rewards.append(-penalty_tracker.penalty) # print("net penalty: %.5f" % penalty_tracker.penalty) # input('...') tar = nvm.registers["tar"] # decode has some robustness to noise even if tar connections are trained target_changed = (tar.decode(tar.content) != tar.decode( tar.old_content)) if done: break if len(rewards) == 0: # target never changed mu = nvm.registers["jnt"].content dist = tr.distributions.normal.Normal(mu, 0.001) log_probs.append(dist.log_prob(mu).sum()) # multivariate white noise rewards = [-10] sym_reward = compute_symbolic_reward(nvm.env, goal_thing_below) spa_reward = compute_spatial_reward(nvm.env, goal_thing_below) end_reward = calc_reward(sym_reward, spa_reward) rewards[-1] += end_reward return end_reward, log_prob, rewards, log_probs
# constant line pt.plot([0, len(avg_rewards)], avg_rewards[[0, 0]], 'g-') pt.show() if showtrained: rep = 0 with open("pac_state_%d.pkl" % rep, "rb") as f: (init_regs, init_conns) = pk.load(f) env = BlocksWorldEnv(show=True, step_hook=penalty_tracker.step_hook) env.load_blocks(thing_below) # set up rvm and virtualize rvm = make_abstract_machine(env, num_bases, max_levels) memorize_env(rvm, goal_thing_above) rvm.reset({"jnt": "rest"}) rvm.mount("main") nvm = virtualize(rvm, σ) nvm.reset_state(init_regs, init_conns) reward, log_prob, rewards, log_probs = run_episode(env, thing_below, goal_thing_below, nvm, init_regs, init_conns, penalty_tracker, sigma=0)