def sample_states(env, q_fn, visitation_probs, n_sample, ent_wt): dS, dA = visitation_probs.shape samples = np.random.choice(np.arange(dS * dA), size=n_sample, p=visitation_probs.reshape(dS * dA)) policy = get_policy(q_fn, ent_wt=ent_wt) observations = samples // dA actions = samples % dA a_logprobs = np.log(policy[observations, actions]) observations_next = [] for i in range(n_sample): t_distr = env.tabular_trans_distr(observations[i], actions[i]) next_state = flat_to_one_hot(np.random.choice(np.arange(len(t_distr)), p=t_distr), ndim=dS) observations_next.append(next_state) observations_next = np.array(observations_next) return { 'observations': flat_to_one_hot(observations, ndim=dS), 'actions': flat_to_one_hot(actions, ndim=dA), 'a_logprobs': a_logprobs, 'observations_next': observations_next }
def compute_visitation(env, q_fn, ent_wt=1.0, T=50, discount=1.0): pol_probs = get_policy(q_fn, ent_wt=ent_wt) dim_obs = env.observation_space.flat_dim dim_act = env.action_space.flat_dim state_visitation = np.expand_dims(env.initial_state_distribution, axis=1) t_matrix = env.transition_matrix # S x A x S sa_visit_t = np.zeros((dim_obs, dim_act, T)) for i in range(T): sa_visit = state_visitation * pol_probs sa_visit_t[:, :, i] = sa_visit #(discount**i) * sa_visit # sum-out (SA)S new_state_visitation = np.einsum('ij,ijk->k', sa_visit, t_matrix) state_visitation = np.expand_dims(new_state_visitation, axis=1) return np.sum(sa_visit_t, axis=2) / float(T)
def tabular_gcl_irl(env, demo_visitations, irl_model, num_itrs=50, ent_wt=1.0, lr=1e-3, state_only=False, discount=0.99, batch_size=20024): dim_obs = env.observation_space.flat_dim dim_act = env.action_space.flat_dim states_all = [] actions_all = [] for s in range(dim_obs): for a in range(dim_act): states_all.append(flat_to_one_hot(s, dim_obs)) actions_all.append(flat_to_one_hot(a, dim_act)) states_all = np.array(states_all) actions_all = np.array(actions_all) path_all = {'observations': states_all, 'actions': actions_all} # Initialize policy and reward function reward_fn = np.zeros((dim_obs, dim_act)) q_rew = np.zeros((dim_obs, dim_act)) update = adam_optimizer(lr) for it in TrainingIterator(num_itrs, heartbeat=1.0): q_itrs = 20 if it.itr > 5 else 100 ### compute policy in closed form q_rew = q_iteration(env, reward_matrix=reward_fn, ent_wt=ent_wt, warmstart_q=q_rew, K=q_itrs, gamma=discount) pol_rew = get_policy(q_rew, ent_wt=ent_wt) ### update reward # need to count how often the policy will visit a particular (s, a) pair pol_visitations = compute_visitation(env, q_rew, ent_wt=ent_wt, T=5, discount=discount) # now we need to sample states and actions, and give them to the discriminator demo_path = sample_states(env, q_rew, demo_visitations, batch_size, ent_wt) irl_model.set_demos([demo_path]) path = sample_states(env, q_rew, pol_visitations, batch_size, ent_wt) irl_model.fit([path], policy=pol_rew, max_itrs=200, lr=1e-3, batch_size=1024) rew_stack = irl_model.eval([path_all])[0] reward_fn = np.zeros_like(q_rew) i = 0 for s in range(dim_obs): for a in range(dim_act): reward_fn[s, a] = rew_stack[i] i += 1 diff_visit = np.abs(demo_visitations - pol_visitations) it.record('VisitationDiffInfNorm', np.max(diff_visit)) it.record('VisitationDiffAvg', np.mean(diff_visit)) if it.heartbeat: print(it.itr_message()) print('\tVisitationDiffInfNorm:', it.pop_mean('VisitationDiffInfNorm')) print('\tVisitationDiffAvg:', it.pop_mean('VisitationDiffAvg')) print('visitations', pol_visitations) print('diff_visit', diff_visit) adjusted_rew = reward_fn - np.mean(reward_fn) + np.mean( env.rew_matrix) print('adjusted_rew', adjusted_rew) return reward_fn, q_rew
env = random_env(16, 4, seed=1, terminate=False, t_sparsity=0.8) env2 = random_env(16, 4, seed=2, terminate=False, t_sparsity=0.8) #plotter = TabularPlotter(4, 16, invert_y=True, text_values=False) dS = env.spec.observation_space.flat_dim dU = env.spec.action_space.flat_dim dO = 8 ent_wt = 0.5 discount = 0.7 obs_matrix = np.random.randn(dS, dO) true_q = q_iteration(env, K=150, ent_wt=ent_wt, gamma=discount) true_sa_visits = compute_visitation(env, true_q, ent_wt=ent_wt, T=5, discount=discount) expert_pol = get_policy(true_q, ent_wt=ent_wt) if True: learned_rew, learned_q = tabular_maxent_irl(env, true_sa_visits, lr=0.01, num_itrs=1000, ent_wt=ent_wt, state_only=False, discount=discount) #extracted_rew = get_reward(env, learned_q, ent_wt=ent_wt, gamma=discount) #new_q = q_iteration(env, K=150, ent_wt=ent_wt, gamma=discount, reward_matrix=extracted_rew) learned_pol = get_policy(learned_q, ent_wt=ent_wt) else: import tensorflow as tf