"""

Code to load a policy and generate rollout data. Adapted from https://github.com/berkeleydeeprlcourse.
Example usage:
    python run_policy.py ../trained_policies/Humanoid-v1/policy_reward_11600/lin_policy_plus.npz Humanoid-v1 --render \
            --num_rollouts 20
"""
import numpy as np
import gym


def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('expert_policy_file', type=str)
    parser.add_argument('envname', type=str)
    parser.add_argument('--render', action='store_true')
    parser.add_argument('--num_rollouts', type=int, default=20,
                        help='Number of expert rollouts')
    args = parser.parse_args()

    print('loading and building expert policy')
    lin_policy = np.load(args.expert_policy_file)
    lin_policy = lin_policy[lin_policy.files[0]]

    M = lin_policy[0]
    # mean and std of state vectors estimated online by ARS.
    mean = lin_policy[1]
    std = lin_policy[2]

    env = gym.make(args.envname)

    returns = []
    observations = []
    actions = []
    for i in range(args.num_rollouts):
        print('iter', i)
        obs = env.reset()
        done = False
        totalr = 0.
        steps = 0
        while not done:
            action = np.dot(M, (obs - mean)/std)
            observations.append(obs)
            actions.append(action)


            obs, r, done, _ = env.step(action)
            totalr += r
            steps += 1
            if args.render:
                env.render()
            if steps % 100 == 0: print("%i/%i"%(steps, env.spec.timestep_limit))
            if steps >= env.spec.timestep_limit:
                break
        returns.append(totalr)

    print('returns', returns)
    print('mean return', np.mean(returns))
    print('std of return', np.std(returns))

if __name__ == '__main__':
    main()

[((14, 13, 14, 38), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((23, 17, 23, 49), 'numpy.load', 'np.load', ({(23, 25, 23, 48): 'args.expert_policy_file'}, {}), '(args.expert_policy_file)', True, 'import numpy as np\n'), ((31, 10, 31, 32), 'gym.make', 'gym.make', ({(31, 19, 31, 31): 'args.envname'}, {}), '(args.envname)', False, 'import gym\n'), ((59, 25, 59, 41), 'numpy.mean', 'np.mean', ({(59, 33, 59, 40): 'returns'}, {}), '(returns)', True, 'import numpy as np\n'), ((60, 27, 60, 42), 'numpy.std', 'np.std', ({(60, 34, 60, 41): 'returns'}, {}), '(returns)', True, 'import numpy as np\n'), ((43, 21, 43, 48), 'numpy.dot', 'np.dot', ({(43, 28, 43, 29): 'M', (43, 31, 43, 47): '(obs - mean) / std'}, {}), '(M, (obs - mean) / std)', True, 'import numpy as np\n')]