test_policy.py
__author__ = 'yuwenhao'
import gym
from baselines.common import set_global_seeds, tf_util as U
from baselines import bench
import os.path as osp
import sys, os, time, errno
import joblib
import numpy as np
import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
from gym import wrappers
import tensorflow as tf
from baselines.ppo1 import mlp_policy, pposgd_simple
import baselines.common.tf_util as U
import pydart2.utils.transformations as trans
import json
np.random.seed(1)
def policy_fn(name, ob_space, ac_space):
return mlp_policy.MlpPolicy(name=name, ob_space=ob_space, ac_space=ac_space,
hid_size=64, num_hid_layers=3, gmm_comp=1)
def save_one_frame_shape(env, fpath, step):
robo_skel = env.env.robot_skeleton
data = []
for b in robo_skel.bodynodes:
if len(b.shapenodes) == 0:
continue
if 'cover' in b.name:
continue
shape_transform = b.T.dot(b.shapenodes[0].relative_transform()).tolist()
#pos = trans.translation_from_matrix(shape_transform)
#rot = trans.euler_from_matrix(shape_transform)
shape_class = str(type(b.shapenodes[0].shape))
if 'Mesh' in shape_class:
stype = 'Mesh'
path = b.shapenodes[0].shape.path()
scale = b.shapenodes[0].shape.scale().tolist()
sub_data = [path, scale]
elif 'Box' in shape_class:
stype = 'Box'
sub_data = b.shapenodes[0].shape.size().tolist()
elif 'Ellipsoid' in shape_class:
stype = 'Ellipsoid'
sub_data = b.shapenodes[0].shape.size().tolist()
elif 'MultiSphere' in shape_class:
stype = 'MultiSphere'
sub_data = b.shapenodes[0].shape.spheres()
for s in range(len(sub_data)):
sub_data[s]['pos'] = sub_data[s]['pos'].tolist()
data.append([stype, b.name, shape_transform, sub_data])
file = fpath + '/frame_' + str(step)+'.txt'
json.dump(data, open(file, 'w'))
if __name__ == '__main__':
save_render_data = False
interpolate = 0
prev_state = None
render_step = 0
render_path = 'render_data/' + 'humanoid_walk'
try:
os.makedirs(render_path)
except OSError as e:
if e.errno != errno.EEXIST:
raise
if len(sys.argv) > 1:
env = gym.make(sys.argv[1])
else:
env = gym.make('DartWalker3d-v1')
if hasattr(env.env, 'disableViewer'):
env.env.disableViewer = False
# manually set the target velocities for different tasks
if len(sys.argv) > 2:
policy_directory = '/'.join(sys.argv[2].split('/')[0:-1])+'/' # put data back to the folder that stores the policies
if sys.argv[2][0] == '/':
policy_directory = '/' + policy_directory
save_directory = policy_directory + '/stats'
try:
os.makedirs(save_directory)
except OSError as e:
if e.errno != errno.EEXIST:
raise
if sys.argv[1] == 'DartWalker3d-v1':
env.env.assist_timeout = 0.0
if 'walk' in sys.argv[2]: # walk task
env.env.final_tv = 1.0
env.env.tv_endtime = 0.5
if 'run' in sys.argv[2]: # run task
env.env.final_tv = 5.0
env.env.tv_endtime = 2.0
if sys.argv[1] == 'DartHumanWalker-v1':
env.env.assist_timeout = 0
if 'walk' in sys.argv[2]: # walk task
env.env.final_tv = 1.5
env.env.tv_endtime = 0.5
if 'walk_back' in sys.argv[2]: # walk back task
env.env.final_tv = -1.5
env.env.tv_endtime = 0.5
if 'run' in sys.argv[2]: # run task
env.env.final_tv = 5.0
env.env.tv_endtime = 3.0
if sys.argv[1] == 'DartDogRobot-v1':
env.env.assist_timeout = 0.0
if 'walk' in sys.argv[2]: # walk task
env.env.final_tv = 2.0
env.env.tv_endtime = 1.0
if 'run' in sys.argv[2]: # run task
env.env.final_tv = 7.0
env.env.tv_endtime = 3.0
if sys.argv[1] == 'DartHexapod-v1':
env.env.assist_timeout = 0.0
if 'walk' in sys.argv[2]: # walk task
env.env.final_tv = 2.0
env.env.tv_endtime = 1.0
if 'run' in sys.argv[2]: # run task
env.env.final_tv = 4.0
env.env.tv_endtime = 2.0
record = False
if len(sys.argv) > 3:
record = int(sys.argv[3]) == 1
if record:
env_wrapper = wrappers.Monitor(env, save_directory, force=True)
else:
env_wrapper = env
if len(sys.argv) > 4:
env.env.visualize = int(sys.argv[4]) == 1
if hasattr(env.env, 'reset_range'):
env.env.reset_range = 0.0
sess = tf.InteractiveSession()
policy = None
if len(sys.argv) > 2:
policy_params = joblib.load(sys.argv[2])
ob_space = env.observation_space
ac_space = env.action_space
policy = policy_fn("pi", ob_space, ac_space)
U.initialize()
cur_scope = policy.get_variables()[0].name[0:policy.get_variables()[0].name.find('/')]
orig_scope = list(policy_params.keys())[0][0:list(policy_params.keys())[0].find('/')]
vars = policy.get_variables()
for i in range(len(policy.get_variables())):
assign_op = policy.get_variables()[i].assign(
policy_params[policy.get_variables()[i].name.replace(cur_scope, orig_scope, 1)])
sess.run(assign_op)
if 'curriculum' in sys.argv[2] and 'policy_params.pkl' in sys.argv[2]:
if os.path.isfile(sys.argv[2].replace('policy_params.pkl', 'init_poses.pkl')):
init_qs, init_dqs = joblib.load(sys.argv[2].replace('policy_params.pkl', 'init_poses.pkl'))
env.env.init_qs = init_qs
env.env.init_dqs = init_dqs
'''ref_policy_params = joblib.load('data/ppo_DartHumanWalker-v1210_energy015_vel65_6s_mirror_up01fwd01ltl15_spinepen1yaw001_thighyawpen005_initbentelbow_velrew3_avg_dcon1_asinput_damping2kneethigh_thigh150knee100_curriculum_1xjoint_shoulder90_dqpen00001/policy_params.pkl')
ref_policy = policy_fn("ref_pi", ob_space, ac_space)
cur_scope = ref_policy.get_variables()[0].name[0:ref_policy.get_variables()[0].name.find('/')]
orig_scope = list(ref_policy_params.keys())[0][0:list(ref_policy_params.keys())[0].find('/')]
vars = ref_policy.get_variables()
for i in range(len(ref_policy.get_variables())):
assign_op = ref_policy.get_variables()[i].assign(
ref_policy_params[ref_policy.get_variables()[i].name.replace(cur_scope, orig_scope, 1)])
sess.run(assign_op)
env.env.ref_policy = ref_policy'''
#init_q, init_dq = joblib.load('data/skel_data/init_states.pkl')
#env.env.init_qs = init_q
#env.env.init_dqs = init_dq
print('===================')
o = env_wrapper.reset()
rew = 0
actions = []
traj = 1
ct = 0
vel_rew = []
action_pen = []
deviation_pen = []
ref_rewards = []
ref_feat_rew = []
rew_seq = []
com_z = []
x_vel = []
foot_contacts = []
contact_force = []
both_contact_forces = []
avg_vels = []
d=False
step = 0
total_steps = 0
save_qs = []
save_dqs = []
save_init_state = False
while ct < traj:
if policy is not None:
ac, vpred = policy.act(step<0, o) # apply stochastic policy at the beginning
act = ac
else:
act = env.action_space.sample()
actions.append(act)
'''if env_wrapper.env.env.t > 3.0 and env_wrapper.env.env.t < 6.0:
env_wrapper.env.env.robot_skeleton.bodynode('head').add_ext_force(np.array([-200, 0, 0]))'''
o, r, d, env_info = env_wrapper.step(act)
if 'action_pen' in env_info:
action_pen.append(env_info['action_pen'])
if 'vel_rew' in env_info:
vel_rew.append(env_info['vel_rew'])
rew_seq.append(r)
if 'deviation_pen' in env_info:
deviation_pen.append(env_info['deviation_pen'])
if 'contact_force' in env_info:
contact_force.append(env_info['contact_force'])
if 'contact_forces' in env_info:
both_contact_forces.append(env_info['contact_forces'])
if 'ref_reward' in env_info:
ref_rewards.append(env_info['ref_reward'])
if 'ref_feat_rew' in env_info:
ref_feat_rew.append(env_info['ref_feat_rew'])
if 'avg_vel' in env_info:
avg_vels.append(env_info['avg_vel'])
com_z.append(o[1])
foot_contacts.append(o[-2:])
rew += r
if len(sys.argv) > 4:
if env.env.visualize:
env_wrapper.render()
else:
env_wrapper.render()
step += 1
total_steps += 1
#time.sleep(0.1)
if len(o) > 25:
x_vel.append(env.env.robot_skeleton.dq[0])
#if np.abs(env.env.t - env.env.tv_endtime) < 0.01:
# save_qs.append(env.env.robot_skeleton.q)
save_dqs.append(env.env.robot_skeleton.dq)
if save_render_data:
cur_state = env.env.state_vector()
if prev_state is not None and interpolate > 0:
for it in range(interpolate):
int_state = (it+1)*1.0/(interpolate+1) * prev_state + (1-(it+1)*1.0/(interpolate+1)) * cur_state
env.env.set_state_vector(int_state)
save_one_frame_shape(env, render_path, render_step)
render_step += 1
env.env.set_state_vector(cur_state)
save_one_frame_shape(env, render_path, render_step)
render_step += 1
prev_state = env.env.state_vector()
if d:
step = 0
if 'contact_locations' in env_info:
c_loc = env_info['contact_locations']
for j in range(len(c_loc[0]) - 1):
c_loc[0][j] = c_loc[0][j+1] - c_loc[0][j]
for j in range(len(c_loc[1]) - 1):
c_loc[1][j] = c_loc[1][j + 1] - c_loc[1][j]
print(np.mean(c_loc[0][0:-1], axis=0))
print(np.mean(c_loc[1][0:-1], axis=0))
ct += 1
print('reward: ', rew)
o=env_wrapper.reset()
#break
print('avg rew ', rew / traj)
print('avg energy penalty: ', np.sum(action_pen)/total_steps)
print('total vel rew: ', np.sum(vel_rew)/traj)
if len(sys.argv) > 2:
np.savetxt(save_directory+'/average_action_magnitude.txt', [np.sum(action_pen)/total_steps])
if 'Walker' in sys.argv[1]: # measure SI for biped
l_contact_total = 0
r_contact_total = 0
for i in range(len(actions)):
l_contact_total += np.linalg.norm(actions[i][[0,1,2,3,4,5]])
r_contact_total += np.linalg.norm(actions[i][[6,7,8,9,10,11]])
print('total forces: ', l_contact_total, r_contact_total)
print('SI: ', 2*np.abs(l_contact_total-r_contact_total)/(l_contact_total+r_contact_total))
np.savetxt(save_directory + '/symmetry_index.txt', [2*np.abs(l_contact_total-r_contact_total)/(l_contact_total+r_contact_total)])
