import logging
logging.disable(logging.CRITICAL)
import math
from tabulate import tabulate
from mjrl.utils.make_train_plots import make_train_plots
from mjrl.utils.gym_env import GymEnv
from mjrl.samplers.core import sample_paths
import numpy as np
import torch
import pickle
import imageio
import time as timer
import os
import copy
import matplotlib.pyplot as plt

try:
    import exptools
    from colorsys import hsv_to_rgb
    import pyvista as pv
except ImportError:
    exptools = None

def _load_latest_policy_and_logs(agent, *, policy_dir, logs_dir):
    """Loads the latest policy.
    Returns the next step number to begin with.
    """
    assert os.path.isdir(policy_dir), str(policy_dir)
    assert os.path.isdir(logs_dir), str(logs_dir)

    log_csv_path = os.path.join(logs_dir, 'log.csv')
    if not os.path.exists(log_csv_path):
        return 0   # fresh start

    print("Reading: {}".format(log_csv_path))
    agent.logger.read_log(log_csv_path)
    last_step = agent.logger.max_len - 1
    if last_step <= 0:
        return 0   # fresh start


    # find latest policy/baseline
    i = last_step
    while i >= 0:
        policy_path = os.path.join(policy_dir, 'policy_{}.pickle'.format(i))
        baseline_path = os.path.join(policy_dir, 'baseline_{}.pickle'.format(i))

        if not os.path.isfile(policy_path):
            i = i -1
            continue
        else:
            print("Loaded last saved iteration: {}".format(i))

        with open(policy_path, 'rb') as fp:
            agent.policy = pickle.load(fp)
        with open(baseline_path, 'rb') as fp:
            agent.baseline = pickle.load(fp)

        # additional
        # global_status_path = os.path.join(policy_dir, 'global_status.pickle')
        # with open(global_status_path, 'rb') as fp:
        #     agent.load_global_status( pickle.load(fp) )

        agent.logger.shrink_to(i + 1)
        assert agent.logger.max_len == i + 1
        return agent.logger.max_len

    # cannot find any saved policy
    raise RuntimeError("Log file exists, but cannot find any saved policy.")

def save_voxel_visualization(obj_name, reset_mode_conf, reward_conf, obj_orientation, obj_relative_position, obj_scale, pc_frame, iternum, is_best_policy):
    uniform_gt_data = np.load("/home/jianrenw/prox/tslam/assets/uniform_gt/uniform_{}_o3d.npz".format(obj_name))['pcd']
    data_scale = uniform_gt_data * obj_scale
    data_rotate = data_scale.copy()

    x = data_rotate[:, 0].copy()
    y = data_rotate[:, 1].copy()
    z = data_rotate[:, 2].copy()
    x_theta = obj_orientation[0]
    data_rotate[:, 0] = x
    data_rotate[:, 1] = y*math.cos(x_theta) - z*math.sin(x_theta)
    data_rotate[:, 2] = y*math.sin(x_theta) + z*math.cos(x_theta)

    x = data_rotate[:, 0].copy()
    y = data_rotate[:, 1].copy()
    z = data_rotate[:, 2].copy()
    y_theta = obj_orientation[1]
    data_rotate[:, 0] = x * math.cos(y_theta) + z * math.sin(y_theta)
    data_rotate[:, 1] = y
    data_rotate[:, 2] = z * math.cos(y_theta) - x * math.sin(y_theta)

    x = data_rotate[:, 0].copy()
    y = data_rotate[:, 1].copy()
    z = data_rotate[:, 2].copy()
    z_theta = obj_orientation[2]
    data_rotate[:, 0] = x * math.cos(z_theta) - y * math.sin(z_theta)
    data_rotate[:, 1] = x * math.sin(z_theta) + y * math.cos(z_theta)
    data_rotate[:, 2] = z

    data_trans = data_rotate.copy()
    data_trans[:, 0] += obj_relative_position[0]
    data_trans[:, 1] += obj_relative_position[1]
    data_trans[:, 2] += obj_relative_position[2]

    uniform_gt_data = data_trans.copy()
    data = pc_frame
    resolution = 0.01
    sep_x = math.ceil(0.3 / resolution)
    sep_y = math.ceil(0.3 / resolution)
    sep_z = math.ceil(0.3 / resolution)
    x, y, z = np.indices((sep_x, sep_y, sep_z))

    cube1 = (x<0) & (y <1) & (z<1)
    gtcube = (x<0) & (y <1) & (z<1)
    voxels = cube1
    gt_voxels = gtcube

    # draw gt
    gt_map_list = []
    for idx,val in enumerate(uniform_gt_data):
        idx_x = math.floor((val[0] + 0.15) / resolution)
        idx_y = math.floor((val[1] + 0.15) / resolution)
        idx_z = math.floor((val[2]) / resolution)
        # if idx_z > 6:
        #     continue
        name = str(idx_x) + '_' + str(idx_y) + '_' + str(idx_z)
        if name not in gt_map_list:
            gt_map_list.append(name)
        cube = (x < idx_x + 1) & (y < idx_y + 1) & (z < idx_z + 1) & (x >= idx_x) & (y >= idx_y) & (z >= idx_z)
        # combine the objects into a single boolean array
        gt_voxels += cube

    # draw cuboids in the top left and bottom right corners, and a link between them
    map_list = []
    for idx,val in enumerate(data):
        idx_x = math.floor((val[0] + 0.15) / resolution)
        idx_y = math.floor((val[1] + 0.15) / resolution)
        idx_z = math.floor((val[2]) / resolution)
        # if idx_z > 6:
        #     continue
        name = str(idx_x) + '_' + str(idx_y) + '_' + str(idx_z)
        if name not in map_list and name in gt_map_list:
            map_list.append(name)
        cube = (x < idx_x + 1) & (y < idx_y + 1) & (z < idx_z + 1) & (x >= idx_x) & (y >= idx_y) & (z >= idx_z)
        # combine the objects into a single boolean array
        voxels += cube

    # gt_obj4:668
    occupancy = len(map_list) / len(gt_map_list)
    # print(len(map_list) / sep_x / sep_y / sep_z )

    is_best_reconstruct = True
    files = os.listdir('/home/jianrenw/prox/tslam/data/result/best_eval/{}/{}/{}/'.format(obj_name, reset_mode_conf, reward_conf))
    for file in files:
        if "overlap" in file and "png" in file:
            file_str = str(file)
            previous_occup = file_str[(file_str.index("-")+1):file_str.index(".png")]
            if occupancy < float(previous_occup):
                is_best_reconstruct = False
    # obj_name = "obj{}".format(obj_name)
    # set the colors of each object
    vis_voxel = gt_voxels | voxels
    colors = np.empty(vis_voxel.shape, dtype=object)
    colors[gt_voxels] = 'white'
    colors[voxels] = 'cyan'
    # and plot everything
    ax = plt.figure().add_subplot(projection='3d')
    ax.set_zlim(1,30)
    ax.voxels(vis_voxel, facecolors=colors, edgecolor='g', alpha=.4, linewidth=.05)
    # plt.savefig('uniform_gtbox_{}.png'.format(step))

    if is_best_policy or is_best_reconstruct:
        plt.savefig('/home/jianrenw/prox/tslam/data/result/best_eval/{}/{}/{}/bp{}_br{}_overlap-{}.png'.format(obj_name, reset_mode_conf, reward_conf, is_best_policy, is_best_reconstruct, occupancy))    
    plt.savefig('voxel/iter-{}-{}-overlap-{}.png'.format(iternum, obj_name, occupancy))
    plt.close()

    ax = plt.figure().add_subplot(projection='3d')
    ax.set_zlim(1,30)
    ax.voxels(gt_voxels, facecolors=colors, edgecolor='g', alpha=.4, linewidth=.05)
    if is_best_policy or is_best_reconstruct:
        plt.savefig('/home/jianrenw/prox/tslam/data/result/best_eval/{}/{}/{}/gt.png'.format(obj_name, reset_mode_conf, reward_conf))    
    plt.savefig('voxel/iter-{}-{}-gt.png'.format(iternum, obj_name))
    plt.close()

    ax = plt.figure().add_subplot(projection='3d')
    ax.set_zlim(1,30)
    ax.voxels(voxels, facecolors=colors, edgecolor='g', alpha=.4, linewidth=.05)
    if is_best_policy or is_best_reconstruct:
        plt.savefig('/home/jianrenw/prox/tslam/data/result/best_eval/{}/{}/{}/bp{}_br{}_exp.png'.format(obj_name, reset_mode_conf, reward_conf, is_best_policy, is_best_reconstruct))    
    plt.savefig('voxel/iter-{}-{}-exp.png'.format(iternum, obj_name))
    plt.close()
    return is_best_reconstruct, occupancy

def train_agent(job_name, agent,
                seed = 0,
                niter = 101,
                gamma = 0.995,
                gae_lambda = None,
                num_cpu = 16,
                sample_mode = 'trajectories',
                horizon= int(150), 
                num_traj = 50,
                num_samples = 50000, # has precedence, used with sample_mode = 'samples'
                save_freq = 10,
                evaluation_rollouts = None,
                plot_keys = ['stoc_pol_mean'],
                env_kwargs= dict(),
                visualize_kwargs= dict(),
                sample_paths_kwargs= dict(),
                ):
    print("num_cpu{}".format(num_cpu))
    np.random.seed(seed)
    if os.path.isdir(job_name) == False:
        os.mkdir(job_name)
    previous_dir = os.getcwd()
    obj_name = env_kwargs["obj_name"]
    reset_mode_conf = env_kwargs["reset_mode"]
    reward_conf = "cf{}knn{}voxel{}".format(env_kwargs["chamfer_r_factor"], env_kwargs["knn_r_factor"], env_kwargs["new_voxel_r_factor"])
    os.chdir(job_name) # important! we are now in the directory to save data
    if os.path.isdir('iterations') == False: os.mkdir('iterations')
    if os.path.isdir('2dpointcloud') == False: os.mkdir('2dpointcloud')
    if os.path.isdir('pointcloudnpz') == False: os.mkdir('pointcloudnpz')
    if os.path.isdir('/home/jianrenw/prox/tslam/data/result/best_policy/{}/{}/{}'.format(obj_name, reset_mode_conf, reward_conf)) == False: os.makedirs('/home/jianrenw/prox/tslam/data/result/best_policy/{}/{}/{}'.format(obj_name, reset_mode_conf, reward_conf))
    if os.path.isdir('/home/jianrenw/prox/tslam/data/result/best_eval/{}/{}/{}'.format(obj_name, reset_mode_conf, reward_conf)) == False: os.makedirs('/home/jianrenw/prox/tslam/data/result/best_eval/{}/{}/{}'.format(obj_name, reset_mode_conf, reward_conf))
    if os.path.isdir('voxel') == False: os.mkdir('voxel')
    if os.path.isdir('logs') == False and agent.save_logs == True: os.mkdir('logs')
    best_policy = copy.deepcopy(agent.policy)
    best_perf = -1e8
    train_curve = best_perf*np.ones(niter)
    mean_pol_perf = 0.0
    e = GymEnv(agent.env.env_id, env_kwargs)

    # Load from any existing checkpoint, policy, statistics, etc.
    # Why no checkpointing.. :(
    i_start = _load_latest_policy_and_logs(agent,
                                           policy_dir='iterations',
                                           logs_dir='logs')
    if i_start:
        print("Resuming from an existing job folder ...")

    for i in range(i_start, niter):
        print("......................................................................................")
        print("ITERATION : %i " % i)
        is_best_policy = False
        if train_curve[i-1] > best_perf:
            if exptools: exptools.logging.logger.log_text("update best_policy")
            best_policy = copy.deepcopy(agent.policy)
            best_perf = train_curve[i-1]
            is_best_policy = True

        N = num_traj if sample_mode == 'trajectories' else num_samples
        stats = agent.train_step(
            N=N, 
            sample_mode=sample_mode,
            horizon= horizon, 
            gamma=gamma,
            gae_lambda=gae_lambda,
            num_cpu=num_cpu,
            env_kwargs= env_kwargs,
            sample_paths_kwargs= sample_paths_kwargs,
        )
        train_curve[i] = stats[0]

        if evaluation_rollouts is not None and evaluation_rollouts > 0:
            print("Performing evaluation rollouts ........")
            eval_paths = sample_paths(
                num_traj=evaluation_rollouts,
                env=e.env_id,
                policy=agent.policy,
                eval_mode=True,
                base_seed=seed,
                num_cpu=num_cpu,
                env_kwargs= env_kwargs,
                **sample_paths_kwargs)
            mean_pol_perf = np.mean([np.sum(path['rewards']) for path in eval_paths])
            if agent.save_logs:
                agent.logger.log_kv('eval_score', mean_pol_perf)
                if exptools: exptools.logging.logger.log_scalar('eval_score', mean_pol_perf, i)
            if exptools:
                env_infos = [path["env_infos"] for path in eval_paths] # a list of dict
                rewards = dict()
                total_points = list()
                if env_infos:
                    # get decomposed reward statistics
                    keys = [k for k in env_infos[0].keys() if "_p" in k[-2:] or "_r" in k[-2:] or "occupancy" in k]
                    for k in keys:
                        rewards[k] = list()
                        for env_info in env_infos:
                            rewards[k].append(env_info[k])
                    for env_info in env_infos:
                        total_points.append(len(env_info["pointcloud"]))
                for k, v in rewards.items():
                    exptools.logging.logger.log_scalar_batch(k, v, i)
                exptools.logging.logger.log_scalar_batch("total_num_points", total_points, i)
            print(">>> finish evaluation rollouts")

        if (i % save_freq == 0 and i > 0):
            if agent.save_logs:
                agent.logger.save_log('logs/')
                make_train_plots(log=agent.logger.log, keys=plot_keys, save_loc='logs/')
            obj_orientation = env_kwargs["obj_orientation"]
            obj_relative_position = env_kwargs["obj_relative_position"]
            obj_scale = env_kwargs["obj_scale"]  
            policy_file = 'policy_%i.pickle' % i
            baseline_file = 'baseline_%i.pickle' % i
            pickle.dump(agent.policy, open('iterations/' + policy_file, 'wb'))
            pickle.dump(agent.baseline, open('iterations/' + baseline_file, 'wb'))
            pickle.dump(best_policy, open('iterations/best_policy.pickle', 'wb'))

            pickle.dump(agent.global_status, open('iterations/global_status.pickle', 'wb'))

            # save videos and pointcloud and reconstruted mesh          
            if exptools:
                video, env_infos = e.visualize_policy_offscreen(
                    policy= agent.policy,
                    **visualize_kwargs,
                ) # (T, C, H, W)
                video_explore, env_infos_explore = e.visualize_policy_explore(
                    policy= agent.policy,
                    **visualize_kwargs,
                ) # (T, C, H, W)
                pc_frame = np.array(env_infos[-1]["pointcloud"] if len(env_infos[-1]["pointcloud"]) > 0 else np.empty((0, 3)))

                # 3d voxel visualization
                is_best_reconstruct, occupancy = save_voxel_visualization(obj_name, reset_mode_conf, reward_conf, obj_orientation, obj_relative_position, obj_scale, pc_frame, i, is_best_policy)
                if is_best_policy or is_best_reconstruct:
                    pickle.dump(best_policy, open('/home/jianrenw/prox/tslam/data/result/best_policy/{}/{}/{}/bp{}_br{}_best_policy.pickle'.format(obj_name, reset_mode_conf, reward_conf, is_best_policy, is_best_reconstruct), 'wb'))
                if is_best_policy or is_best_reconstruct:
                    np.savez_compressed("pointcloudnpz/alpha_pointcloud_"+str(i)+".npz",pcd=pc_frame)
                    np.savez_compressed("/home/jianrenw/prox/tslam/data/result/best_eval/{}/{}/{}/bp{}_br{}_alpha_pointcloud_overlap-{}.npz".format(obj_name, reset_mode_conf, reward_conf, is_best_policy, is_best_reconstruct, occupancy), pcd=pc_frame)
                # else:
                #     np.savez_compressed("pointcloudnpz/pointcloud_"+str(i)+".npz",pcd=pc_frame)
                
                # pc_frames.append(pc_frame)
                ax = plt.axes()
                ax.scatter(pc_frame[:, 0], pc_frame[:, 1], cmap='viridis', linewidth=0.5)
                if is_best_policy or is_best_reconstruct:
                    plt.savefig("2dpointcloud/alpha_{}.png".format('2dpointcloud' + str(i)))
                    plt.savefig("/home/jianrenw/prox/tslam/data/result/best_eval/{}/{}/{}/bp{}_br{}_alpha_2dpointcloud_overlap-{}.png".format(obj_name, reset_mode_conf, reward_conf, is_best_policy, is_best_reconstruct, occupancy))
                # else:
                #     plt.savefig("2dpointcloud/{}.png".format('2dpointcloud' + str(i)))
                plt.close()
                # =======================================================
                # if obj_name in ["airplane", "apple", "glass", "cup"]:
                exptools.logging.logger.record_image("rendered", video[-1], i)
                exptools.logging.logger.record_gif("rendered", video, i)
                # exptools.logging.logger.record_image("rendered_explore", video_explore[-1], i)
                # exptools.logging.logger.record_gif("rendered_explore", video_explore, i)

        # print results to console
        if i == 0:
            result_file = open('results.txt', 'w')
            print("Iter | Stoc Pol | Mean Pol | Best (Stoc) \n")
            result_file.write("Iter | Sampling Pol | Evaluation Pol | Best (Sampled) \n")
            result_file.close()
        result_file = open('results.txt', 'a')
        result_file.write("%4i %5.2f %5.2f %5.2f \n" % (i, train_curve[i], mean_pol_perf, best_perf))
        result_file.close()
        if agent.save_logs:
            print_data = sorted(filter(lambda v: np.asarray(v[1]).size == 1,
                                       agent.logger.get_current_log().items()))
            print(tabulate(print_data))
        if exptools:
            exptools.logging.logger.log_scalar("Iter", i, i)
            exptools.logging.logger.log_scalar("SamplingPol", train_curve[i], i)
            exptools.logging.logger.log_scalar("EvaluationPol", mean_pol_perf, i)
            exptools.logging.logger.log_scalar("BestSampled", best_perf, i)
            exptools.logging.logger.dump_data()

    # final save
    pickle.dump(best_policy, open('iterations/best_policy.pickle', 'wb'))
    if agent.save_logs:
        agent.logger.save_log('logs/')
        make_train_plots(log=agent.logger.log, keys=plot_keys, save_loc='logs/')
    os.chdir(previous_dir)

[((2, 0, 2, 33), 'logging.disable', 'logging.disable', ({(2, 16, 2, 32): 'logging.CRITICAL'}, {}), '(logging.CRITICAL)', False, 'import logging\n'), ((28, 11, 28, 36), 'os.path.isdir', 'os.path.isdir', ({(28, 25, 28, 35): 'policy_dir'}, {}), '(policy_dir)', False, 'import os\n'), ((29, 11, 29, 34), 'os.path.isdir', 'os.path.isdir', ({(29, 25, 29, 33): 'logs_dir'}, {}), '(logs_dir)', False, 'import os\n'), ((31, 19, 31, 52), 'os.path.join', 'os.path.join', ({(31, 32, 31, 40): 'logs_dir', (31, 42, 31, 51): '"""log.csv"""'}, {}), "(logs_dir, 'log.csv')", False, 'import os\n'), ((108, 12, 108, 39), 'math.ceil', 'math.ceil', ({(108, 22, 108, 38): '0.3 / resolution'}, {}), '(0.3 / resolution)', False, 'import math\n'), ((109, 12, 109, 39), 'math.ceil', 'math.ceil', ({(109, 22, 109, 38): '0.3 / resolution'}, {}), '(0.3 / resolution)', False, 'import math\n'), ((110, 12, 110, 39), 'math.ceil', 'math.ceil', ({(110, 22, 110, 38): '0.3 / resolution'}, {}), '(0.3 / resolution)', False, 'import math\n'), ((111, 14, 111, 47), 'numpy.indices', 'np.indices', ({(111, 25, 111, 46): '(sep_x, sep_y, sep_z)'}, {}), '((sep_x, sep_y, sep_z))', True, 'import numpy as np\n'), ((163, 13, 163, 52), 'numpy.empty', 'np.empty', (), '', True, 'import numpy as np\n'), ((175, 4, 175, 15), 'matplotlib.pyplot.close', 'plt.close', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((183, 4, 183, 15), 'matplotlib.pyplot.close', 'plt.close', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((191, 4, 191, 15), 'matplotlib.pyplot.close', 'plt.close', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((212, 4, 212, 24), 'numpy.random.seed', 'np.random.seed', ({(212, 19, 212, 23): 'seed'}, {}), '(seed)', True, 'import numpy as np\n'), ((215, 19, 215, 30), 'os.getcwd', 'os.getcwd', ({}, {}), '()', False, 'import os\n'), ((219, 4, 219, 22), 'os.chdir', 'os.chdir', ({(219, 13, 219, 21): 'job_name'}, {}), '(job_name)', False, 'import os\n'), ((227, 18, 227, 45), 'copy.deepcopy', 'copy.deepcopy', ({(227, 32, 227, 44): 'agent.policy'}, {}), '(agent.policy)', False, 'import copy\n'), ((231, 8, 231, 44), 'mjrl.utils.gym_env.GymEnv', 'GymEnv', ({(231, 15, 231, 31): 'agent.env.env_id', (231, 33, 231, 43): 'env_kwargs'}, {}), '(agent.env.env_id, env_kwargs)', False, 'from mjrl.utils.gym_env import GymEnv\n'), ((375, 4, 375, 26), 'os.chdir', 'os.chdir', ({(375, 13, 375, 25): 'previous_dir'}, {}), '(previous_dir)', False, 'import os\n'), ((32, 11, 32, 39), 'os.path.exists', 'os.path.exists', ({(32, 26, 32, 38): 'log_csv_path'}, {}), '(log_csv_path)', False, 'import os\n'), ((121, 16, 121, 56), 'math.floor', 'math.floor', ({(121, 27, 121, 55): '(val[0] + 0.15) / resolution'}, {}), '((val[0] + 0.15) / resolution)', False, 'import math\n'), ((122, 16, 122, 56), 'math.floor', 'math.floor', ({(122, 27, 122, 55): '(val[1] + 0.15) / resolution'}, {}), '((val[1] + 0.15) / resolution)', False, 'import math\n'), ((123, 16, 123, 49), 'math.floor', 'math.floor', ({(123, 27, 123, 48): 'val[2] / resolution'}, {}), '(val[2] / resolution)', False, 'import math\n'), ((136, 16, 136, 56), 'math.floor', 'math.floor', ({(136, 27, 136, 55): '(val[0] + 0.15) / resolution'}, {}), '((val[0] + 0.15) / resolution)', False, 'import math\n'), ((137, 16, 137, 56), 'math.floor', 'math.floor', ({(137, 27, 137, 55): '(val[1] + 0.15) / resolution'}, {}), '((val[1] + 0.15) / resolution)', False, 'import math\n'), ((138, 16, 138, 49), 'math.floor', 'math.floor', ({(138, 27, 138, 48): 'val[2] / resolution'}, {}), '(val[2] / resolution)', False, 'import math\n'), ((213, 7, 213, 30), 'os.path.isdir', 'os.path.isdir', ({(213, 21, 213, 29): 'job_name'}, {}), '(job_name)', False, 'import os\n'), ((214, 8, 214, 26), 'os.mkdir', 'os.mkdir', ({(214, 17, 214, 25): 'job_name'}, {}), '(job_name)', False, 'import os\n'), ((220, 7, 220, 34), 'os.path.isdir', 'os.path.isdir', ({(220, 21, 220, 33): '"""iterations"""'}, {}), "('iterations')", False, 'import os\n'), ((220, 45, 220, 67), 'os.mkdir', 'os.mkdir', ({(220, 54, 220, 66): '"""iterations"""'}, {}), "('iterations')", False, 'import os\n'), ((221, 7, 221, 36), 'os.path.isdir', 'os.path.isdir', ({(221, 21, 221, 35): '"""2dpointcloud"""'}, {}), "('2dpointcloud')", False, 'import os\n'), ((221, 47, 221, 71), 'os.mkdir', 'os.mkdir', ({(221, 56, 221, 70): '"""2dpointcloud"""'}, {}), "('2dpointcloud')", False, 'import os\n'), ((222, 7, 222, 37), 'os.path.isdir', 'os.path.isdir', ({(222, 21, 222, 36): '"""pointcloudnpz"""'}, {}), "('pointcloudnpz')", False, 'import os\n'), ((222, 48, 222, 73), 'os.mkdir', 'os.mkdir', ({(222, 57, 222, 72): '"""pointcloudnpz"""'}, {}), "('pointcloudnpz')", False, 'import os\n'), ((225, 7, 225, 29), 'os.path.isdir', 'os.path.isdir', ({(225, 21, 225, 28): '"""voxel"""'}, {}), "('voxel')", False, 'import os\n'), ((225, 40, 225, 57), 'os.mkdir', 'os.mkdir', ({(225, 49, 225, 56): '"""voxel"""'}, {}), "('voxel')", False, 'import os\n'), ((226, 67, 226, 83), 'os.mkdir', 'os.mkdir', ({(226, 76, 226, 82): '"""logs"""'}, {}), "('logs')", False, 'import os\n'), ((229, 28, 229, 42), 'numpy.ones', 'np.ones', ({(229, 36, 229, 41): 'niter'}, {}), '(niter)', True, 'import numpy as np\n'), ((374, 8, 374, 80), 'mjrl.utils.make_train_plots.make_train_plots', 'make_train_plots', (), '', False, 'from mjrl.utils.make_train_plots import make_train_plots\n'), ((48, 15, 48, 42), 'os.path.isfile', 'os.path.isfile', ({(48, 30, 48, 41): 'policy_path'}, {}), '(policy_path)', False, 'import os\n'), ((55, 27, 55, 42), 'pickle.load', 'pickle.load', ({(55, 39, 55, 41): 'fp'}, {}), '(fp)', False, 'import pickle\n'), ((57, 29, 57, 44), 'pickle.load', 'pickle.load', ({(57, 41, 57, 43): 'fp'}, {}), '(fp)', False, 'import pickle\n'), ((81, 26, 81, 43), 'math.cos', 'math.cos', ({(81, 35, 81, 42): 'x_theta'}, {}), '(x_theta)', False, 'import math\n'), ((81, 48, 81, 65), 'math.sin', 'math.sin', ({(81, 57, 81, 64): 'x_theta'}, {}), '(x_theta)', False, 'import math\n'), ((82, 26, 82, 43), 'math.sin', 'math.sin', ({(82, 35, 82, 42): 'x_theta'}, {}), '(x_theta)', False, 'import math\n'), ((82, 48, 82, 65), 'math.cos', 'math.cos', ({(82, 57, 82, 64): 'x_theta'}, {}), '(x_theta)', False, 'import math\n'), ((88, 28, 88, 45), 'math.cos', 'math.cos', ({(88, 37, 88, 44): 'y_theta'}, {}), '(y_theta)', False, 'import math\n'), ((88, 52, 88, 69), 'math.sin', 'math.sin', ({(88, 61, 88, 68): 'y_theta'}, {}), '(y_theta)', False, 'import math\n'), ((90, 28, 90, 45), 'math.cos', 'math.cos', ({(90, 37, 90, 44): 'y_theta'}, {}), '(y_theta)', False, 'import math\n'), ((90, 52, 90, 69), 'math.sin', 'math.sin', ({(90, 61, 90, 68): 'y_theta'}, {}), '(y_theta)', False, 'import math\n'), ((96, 28, 96, 45), 'math.cos', 'math.cos', ({(96, 37, 96, 44): 'z_theta'}, {}), '(z_theta)', False, 'import math\n'), ((96, 52, 96, 69), 'math.sin', 'math.sin', ({(96, 61, 96, 68): 'z_theta'}, {}), '(z_theta)', False, 'import math\n'), ((97, 28, 97, 45), 'math.sin', 'math.sin', ({(97, 37, 97, 44): 'z_theta'}, {}), '(z_theta)', False, 'import math\n'), ((97, 52, 97, 69), 'math.cos', 'math.cos', ({(97, 61, 97, 68): 'z_theta'}, {}), '(z_theta)', False, 'import math\n'), ((167, 9, 167, 21), 'matplotlib.pyplot.figure', 'plt.figure', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((177, 9, 177, 21), 'matplotlib.pyplot.figure', 'plt.figure', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((185, 9, 185, 21), 'matplotlib.pyplot.figure', 'plt.figure', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((226, 7, 226, 28), 'os.path.isdir', 'os.path.isdir', ({(226, 21, 226, 27): '"""logs"""'}, {}), "('logs')", False, 'import os\n'), ((247, 26, 247, 53), 'copy.deepcopy', 'copy.deepcopy', ({(247, 40, 247, 52): 'agent.policy'}, {}), '(agent.policy)', False, 'import copy\n'), ((266, 25, 274, 38), 'mjrl.samplers.core.sample_paths', 'sample_paths', (), '', False, 'from mjrl.samplers.core import sample_paths\n'), ((364, 12, 364, 60), 'exptools.logging.logger.log_scalar', 'exptools.logging.logger.log_scalar', ({(364, 47, 364, 53): '"""Iter"""', (364, 55, 364, 56): 'i', (364, 58, 364, 59): 'i'}, {}), "('Iter', i, i)", False, 'import exptools\n'), ((365, 12, 365, 80), 'exptools.logging.logger.log_scalar', 'exptools.logging.logger.log_scalar', ({(365, 47, 365, 60): '"""SamplingPol"""', (365, 62, 365, 76): 'train_curve[i]', (365, 78, 365, 79): 'i'}, {}), "('SamplingPol', train_curve[i], i)", False, 'import exptools\n'), ((366, 12, 366, 81), 'exptools.logging.logger.log_scalar', 'exptools.logging.logger.log_scalar', ({(366, 47, 366, 62): '"""EvaluationPol"""', (366, 64, 366, 77): 'mean_pol_perf', (366, 79, 366, 80): 'i'}, {}), "('EvaluationPol', mean_pol_perf, i)", False, 'import exptools\n'), ((367, 12, 367, 75), 'exptools.logging.logger.log_scalar', 'exptools.logging.logger.log_scalar', ({(367, 47, 367, 60): '"""BestSampled"""', (367, 62, 367, 71): 'best_perf', (367, 73, 367, 74): 'i'}, {}), "('BestSampled', best_perf, i)", False, 'import exptools\n'), ((368, 12, 368, 47), 'exptools.logging.logger.dump_data', 'exptools.logging.logger.dump_data', ({}, {}), '()', False, 'import exptools\n'), ((246, 25, 246, 79), 'exptools.logging.logger.log_text', 'exptools.logging.logger.log_text', ({(246, 58, 246, 78): '"""update best_policy"""'}, {}), "('update best_policy')", False, 'import exptools\n'), ((294, 16, 294, 93), 'exptools.logging.logger.log_scalar_batch', 'exptools.logging.logger.log_scalar_batch', ({(294, 57, 294, 75): '"""total_num_points"""', (294, 77, 294, 89): 'total_points', (294, 91, 294, 92): 'i'}, {}), "('total_num_points', total_points, i)", False, 'import exptools\n'), ((300, 16, 300, 88), 'mjrl.utils.make_train_plots.make_train_plots', 'make_train_plots', (), '', False, 'from mjrl.utils.make_train_plots import make_train_plots\n'), ((335, 21, 335, 31), 'matplotlib.pyplot.axes', 'plt.axes', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((342, 16, 342, 27), 'matplotlib.pyplot.close', 'plt.close', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((345, 16, 345, 78), 'exptools.logging.logger.record_image', 'exptools.logging.logger.record_image', ({(345, 53, 345, 63): '"""rendered"""', (345, 65, 345, 74): 'video[-1]', (345, 76, 345, 77): 'i'}, {}), "('rendered', video[-1], i)", False, 'import exptools\n'), ((346, 16, 346, 72), 'exptools.logging.logger.record_gif', 'exptools.logging.logger.record_gif', ({(346, 51, 346, 61): '"""rendered"""', (346, 63, 346, 68): 'video', (346, 70, 346, 71): 'i'}, {}), "('rendered', video, i)", False, 'import exptools\n'), ((362, 18, 362, 38), 'tabulate.tabulate', 'tabulate', ({(362, 27, 362, 37): 'print_data'}, {}), '(print_data)', False, 'from tabulate import tabulate\n'), ((275, 37, 275, 60), 'numpy.sum', 'np.sum', ({(275, 44, 275, 59): "path['rewards']"}, {}), "(path['rewards'])", True, 'import numpy as np\n'), ((278, 29, 278, 95), 'exptools.logging.logger.log_scalar', 'exptools.logging.logger.log_scalar', ({(278, 64, 278, 76): '"""eval_score"""', (278, 78, 278, 91): 'mean_pol_perf', (278, 93, 278, 94): 'i'}, {}), "('eval_score', mean_pol_perf, i)", False, 'import exptools\n'), ((293, 20, 293, 69), 'exptools.logging.logger.log_scalar_batch', 'exptools.logging.logger.log_scalar_batch', ({(293, 61, 293, 62): 'k', (293, 64, 293, 65): 'v', (293, 67, 293, 68): 'i'}, {}), '(k, v, i)', False, 'import exptools\n'), ((322, 109, 322, 125), 'numpy.empty', 'np.empty', ({(322, 118, 322, 124): '(0, 3)'}, {}), '((0, 3))', True, 'import numpy as np\n'), ((360, 49, 360, 65), 'numpy.asarray', 'np.asarray', ({(360, 60, 360, 64): 'v[1]'}, {}), '(v[1])', True, 'import numpy as np\n')]