#!/usr/bin/env python
import chainer

from algs import trpo
from env_makers import EnvMaker
from models import GaussianMLPPolicy, MLPBaseline
from utils import SnapshotSaver
import numpy as np
import os
import logger

log_dir = "data/local/trpo-pendulum"

np.random.seed(42)

# Clean up existing logs
os.system("rm -rf {}".format(log_dir))

with logger.session(log_dir):
    env_maker = EnvMaker('Pendulum-v0')
    env = env_maker.make()
    policy = GaussianMLPPolicy(
        observation_space=env.observation_space,
        action_space=env.action_space,
        env_spec=env.spec,
        hidden_sizes=(64, 64),
        hidden_nonlinearity=chainer.functions.tanh,
    )
    baseline = MLPBaseline(
        observation_space=env.observation_space,
        action_space=env.action_space,
        env_spec=env.spec,
        hidden_sizes=(64, 64),
        hidden_nonlinearity=chainer.functions.tanh,
    )
    trpo(
        env=env,
        env_maker=env_maker,
        n_envs=16,
        policy=policy,
        baseline=baseline,
        batch_size=10000,
        n_iters=100,
        snapshot_saver=SnapshotSaver(log_dir),
    )

[((14, 0, 14, 18), 'numpy.random.seed', 'np.random.seed', ({(14, 15, 14, 17): '(42)'}, {}), '(42)', True, 'import numpy as np\n'), ((19, 5, 19, 28), 'logger.session', 'logger.session', ({(19, 20, 19, 27): 'log_dir'}, {}), '(log_dir)', False, 'import logger\n'), ((20, 16, 20, 39), 'env_makers.EnvMaker', 'EnvMaker', ({(20, 25, 20, 38): '"""Pendulum-v0"""'}, {}), "('Pendulum-v0')", False, 'from env_makers import EnvMaker\n'), ((22, 13, 28, 5), 'models.GaussianMLPPolicy', 'GaussianMLPPolicy', (), '', False, 'from models import GaussianMLPPolicy, MLPBaseline\n'), ((29, 15, 35, 5), 'models.MLPBaseline', 'MLPBaseline', (), '', False, 'from models import GaussianMLPPolicy, MLPBaseline\n'), ((44, 23, 44, 45), 'utils.SnapshotSaver', 'SnapshotSaver', ({(44, 37, 44, 44): 'log_dir'}, {}), '(log_dir)', False, 'from utils import SnapshotSaver\n')]