from typing import DefaultDict
import numpy as np
import torch
import argparse
import os
import math
import gym
import sys
import random
import time
import json
from collections import defaultdict
from pathlib import Path

# import dmc2gym
import copy

from tqdm.auto import tqdm
from rich import print
import utils
from logger import Logger
from video import VideoRecorder

from curl_sac import CurlSacAgent
from torchvision import transforms
import apple_gym.env
from diy_gym.utils import flatten, unflatten


def parse_args():
    parser = argparse.ArgumentParser()
    # environment
    parser.add_argument("--domain_name", default="ApplePick-v0")
    parser.add_argument("--pre_transform_image_size", default=124, type=int)

    parser.add_argument("--image_size", default=84, type=int)
    # parser.add_argument("--action_repeat", default=1, type=int)
    parser.add_argument("--frame_stack", default=3, type=int)
    parser.add_argument("--render", action="store_true")
    # replay buffer
    parser.add_argument("--replay_buffer_capacity", default=50000, type=int)
    # train
    parser.add_argument("--agent", default="curl_sac", type=str)
    parser.add_argument("--init_steps", default=10000, type=int)
    parser.add_argument("--num_train_steps", default=3000000, type=int)
    parser.add_argument("--batch_size", default=32, type=int)
    parser.add_argument("--hidden_dim", default=1024, type=int)
    # eval
    parser.add_argument("--eval_freq", default=2000, type=int)
    parser.add_argument("--num_eval_episodes", default=4, type=int)
    # critic
    parser.add_argument("--critic_lr", default=1e-3, type=float)
    parser.add_argument("--critic_beta", default=0.9, type=float)
    parser.add_argument("--critic_tau", default=0.01, type=float)  # try 0.05 or 0.1
    parser.add_argument(
        "--critic_target_update_freq", default=2, type=int
    )  # try to change it to 1 and retain 0.01 above
    # actor
    parser.add_argument("--actor_lr", default=1e-3, type=float)
    parser.add_argument("--actor_beta", default=0.9, type=float)
    parser.add_argument("--actor_log_std_min", default=-10, type=float)
    parser.add_argument("--actor_log_std_max", default=2, type=float)
    parser.add_argument("--actor_update_freq", default=2, type=int)
    # encoder
    parser.add_argument("--encoder_type", default="mixed", type=str)
    parser.add_argument("--encoder_feature_dim", default=50, type=int)
    parser.add_argument("--encoder_lr", default=1e-3, type=float)
    parser.add_argument("--encoder_tau", default=0.05, type=float)
    parser.add_argument("--num_layers", default=4, type=int)
    parser.add_argument("--num_filters", default=32, type=int)
    parser.add_argument("--curl_latent_dim", default=128, type=int)
    # sac
    parser.add_argument("--discount", default=0.99, type=float)
    parser.add_argument("--init_temperature", default=0.1, type=float)
    parser.add_argument("--alpha_lr", default=1e-4, type=float)
    parser.add_argument("--alpha_beta", default=0.5, type=float)
    # misc
    parser.add_argument("--seed", default=-1, type=int)
    parser.add_argument("--work_dir", default="./runs", type=str)
    parser.add_argument("--save_tb", default=False, action="store_true")
    parser.add_argument("--save_buffer", default=False, action="store_true")
    parser.add_argument("--save_video", default=False, action="store_true")
    parser.add_argument("--save_model", default=False, action="store_true")
    parser.add_argument("--detach_encoder", default=False, action="store_true")
    parser.add_argument("--load", type=str)

    parser.add_argument("--log_interval", default=100, type=int)
    args = parser.parse_args()
    return args

keys_to_monitor=[
    'env_reward/apple_pick/tree/min_fruit_dist_reward',
    'env_reward/apple_pick/tree/gripping_fruit_reward',
    # 'env_reward/apple_pick/tree/force_tree_reward',
    # 'env_reward/apple_pick/tree/force_fruit_reward',
    'env_obs/apple_pick/tree/picks'
                ]

def evaluate(env, agent, video, num_episodes, L, step, args):
    all_ep_rewards = []

    def run_eval_loop(sample_stochastically=True):
        start_time = time.time()
        prefix = "stochastic_" if sample_stochastically else ""
        for i in tqdm(range(num_episodes), desc='eval', unit='ep'):
            obs = env.reset()
            video.init(enabled=(i == 0))
            done = False
            episode_reward = 0
            episode_info = defaultdict(int)
            while not done:
                # center crop image
                if args.encoder_type == "mixed":
                    state, img = utils.split_obs(obs)
                    img = utils.center_crop_image(img, args.image_size)
                    obs = utils.combine_obs(state, img)
                with utils.eval_mode(agent):
                    if sample_stochastically:
                        action = agent.sample_action(obs)
                    else:
                        action = agent.select_action(obs)
                obs, reward, done, info = env.step(action)

                for k in keys_to_monitor:
                    episode_info[k] += info[k]
                video.record(env, yaw=i)
                episode_reward += reward

            for k in keys_to_monitor:
                L.log("eval/" + prefix + k, np.sum(episode_info[k]), step)
            video.save("%d.mp4" % step)
            L.log("eval/" + prefix + "episode_reward", episode_reward, step)
            all_ep_rewards.append(episode_reward)

        L.log("eval/" + prefix + "eval_time", time.time() - start_time, step)
        mean_ep_reward = np.mean(all_ep_rewards)
        best_ep_reward = np.max(all_ep_rewards)
        L.log("eval/" + prefix + "mean_episode_reward", mean_ep_reward, step)
        L.log("eval/" + prefix + "best_episode_reward", best_ep_reward, step)

    run_eval_loop(sample_stochastically=False)
    L.dump(step)


def make_agent(obs_shape, action_shape, args, device):
    if args.agent == "curl_sac":
        return CurlSacAgent(
            obs_shape=obs_shape,
            action_shape=action_shape,
            device=device,
            hidden_dim=args.hidden_dim,
            discount=args.discount,
            init_temperature=args.init_temperature,
            alpha_lr=args.alpha_lr,
            alpha_beta=args.alpha_beta,
            actor_lr=args.actor_lr,
            actor_beta=args.actor_beta,
            actor_log_std_min=args.actor_log_std_min,
            actor_log_std_max=args.actor_log_std_max,
            actor_update_freq=args.actor_update_freq,
            critic_lr=args.critic_lr,
            critic_beta=args.critic_beta,
            critic_tau=args.critic_tau,
            critic_target_update_freq=args.critic_target_update_freq,
            encoder_type=args.encoder_type,
            encoder_feature_dim=args.encoder_feature_dim,
            encoder_lr=args.encoder_lr,
            encoder_tau=args.encoder_tau,
            num_layers=args.num_layers,
            num_filters=args.num_filters,
            log_interval=args.log_interval,
            detach_encoder=args.detach_encoder,
            curl_latent_dim=args.curl_latent_dim,
        )
    else:
        assert "agent is not supported: %s" % args.agent


def main():
    import logging
    from rich.logging import RichHandler
    logging.basicConfig(level=logging.INFO, handlers=[RichHandler(rich_tracebacks=True, markup=True)])

    args = parse_args()
    if args.seed == -1:
        args.__dict__["seed"] = np.random.randint(1, 1000000)
        print('seed', args.__dict__["seed"])
    print(args)
    utils.set_seed_everywhere(args.seed)

    env = gym.make(args.domain_name, render=args.render)
    print(env)
    # TODO action repeat wrapper?

    env.seed(args.seed)

    # # stack several consecutive frames together
    if args.encoder_type == "mixed":
        from apple_gym.env.wrappers import FrameStack, ImageState, PermuteImages

        env = FrameStack(
            PermuteImages(ImageState(env), keys=["img"]),
            n=args.frame_stack,
            keys=["img"],
        )

    if args.load =='auto':
        load_dirs = Path(args.work_dir).glob('*/model/curl*.pt')
        load_dirs = sorted(set([str(d.parent) for d in load_dirs]))
        print('load_dirs', load_dirs)
        args.load = str(load_dirs[-1])
        print('auto load', load_dirs)

    # make directory
    ts = time.gmtime()
    ts = time.strftime("%m-%d", ts)
    env_name = args.domain_name
    exp_name = (
        env_name
        + "-"
        + ts
        + "-im"
        + str(args.image_size)
        + "-b"
        + str(args.batch_size)
        + "-s"
        + str(args.seed)
        + "-"
        + args.encoder_type
    )
    args.work_dir = args.work_dir + "/" + exp_name
    print('work_dir', args.work_dir)

    utils.make_dir(args.work_dir)
    video_dir = utils.make_dir(os.path.join(args.work_dir, "video"))
    model_dir = utils.make_dir(os.path.join(args.work_dir, "model"))
    buffer_dir = utils.make_dir(os.path.join(args.work_dir, "buffer"))

    video = VideoRecorder(video_dir if args.save_video else None)

    with open(os.path.join(args.work_dir, "args.json"), "w") as f:
        json.dump(vars(args), f, sort_keys=True, indent=4)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f'device {device}')

    # shapes
    action_shape = env.action_space.shape
    img = env.observation_space.sample()["img"]
    img_aug = utils.center_crop_image(img, args.image_size)
    obs_shape = {"img": img_aug.shape, "state": env.observation_space["state"].shape}

    replay_buffer = utils.ReplayBuffer(
        obs_space=env.observation_space,
        action_space=env.action_space,
        capacity=args.replay_buffer_capacity,
        batch_size=args.batch_size,
        device=device,
        image_size=args.image_size,
    )

    agent = make_agent(
        obs_shape=obs_shape, action_shape=action_shape, args=args, device=device
    )
    if args.load is not None:
        agent.load_curl(args.load)
    

    # summarize
    obs = env.observation_space.sample()
    state, img = utils.split_obs(obs)
    img_crop = utils.center_crop_image(img, agent.image_size)
    obs_crop = utils.combine_obs(state, img_crop)
    obs_crop['img'] = torch.FloatTensor(obs_crop['img']).to(agent.device).unsqueeze(0)
    obs_crop['state'] = torch.FloatTensor(obs_crop['state']).to(agent.device).unsqueeze(0)
    action = agent.sample_action(obs)
    action = torch.FloatTensor(action).to(agent.device).unsqueeze(0)
    from torchsummaryX import summary
    with torch.no_grad():
        print(agent.critic)
        summary(agent.critic, obs_crop, action)
        print(agent.actor)
        summary(agent.actor, obs_crop)

    L = Logger(args.work_dir, use_tb=args.save_tb)

    episode, episode_reward, done = 0, 0, True
    episode_info = defaultdict(int)
    start_time = time.time()

    for step in tqdm(range(args.num_train_steps), desc="train", unit="step", mininterval=360):
        # evaluate agent periodically

        if (step % args.eval_freq == 0) and (step >= args.eval_freq):
            L.log("eval/episode", episode, step)
            evaluate(env, agent, video, args.num_eval_episodes, L, step, args)
            if args.save_model:
                agent.save_curl(model_dir, step)
            if args.save_buffer:
                replay_buffer.save(buffer_dir)

        if done:
            if step > 0:
                if step % args.log_interval == 0:
                    L.log("train/duration", time.time() - start_time, step)
                    L.dump(step)
                start_time = time.time()
            if step % args.log_interval == 0:
                L.log("train/episode_reward", episode_reward, step)

            for k in keys_to_monitor:
                L.log("train/episode_info" + k, np.sum(episode_info[k]), step)

            obs = env.reset()
            assert env.observation_space.contains(
                obs
            ), f"obs should be in space. ob={obs} space={env.observation_space}"
            done = False
            episode_reward = 0
            episode_info = defaultdict(int)
            episode_step = 0
            episode += 1
            if step % args.log_interval == 0:
                L.log("train/episode", episode, step)

        # sample action for data collection
        if step < args.init_steps:
            action = env.action_space.sample()
        else:
            with utils.eval_mode(agent):
                action = agent.sample_action(obs)
        assert env.action_space.contains(
            action
        ), f"obs should be in space. ob={action} space={env.action_space}"

        if step % 10 ==0:
            # run training update
            if step >= args.init_steps:
                num_updates = 1
                for _ in range(num_updates):
                    agent.update(replay_buffer, L, step)

        next_obs, reward, done, info = env.step(action)

        # allow infinite bootstrap
        done_bool = 0 if episode_step + 1 == env._max_episode_steps else float(done)
        episode_reward += reward
        replay_buffer.add(obs, action, reward, next_obs, done_bool)
        for k in keys_to_monitor:
            episode_info[k] += info[k]

        obs = next_obs
        episode_step += 1


if __name__ == "__main__":
    torch.multiprocessing.set_start_method("spawn")

    main()

from typing import DefaultDict
import numpy as np
import torch
import argparse
import os
import math
import gym
import sys
import random
import time
import json
from collections import defaultdict
from pathlib import Path

# import dmc2gym
import copy

from tqdm.auto import tqdm
from rich import print
import utils
from logger import Logger
from video import VideoRecorder

from curl_sac import CurlSacAgent
from torchvision import transforms
import apple_gym.env
from diy_gym.utils import flatten, unflatten


def parse_args():
    parser = argparse.ArgumentParser()
    # environment
    parser.add_argument("--domain_name", default="ApplePick-v0")
    parser.add_argument("--pre_transform_image_size", default=124, type=int)

    parser.add_argument("--image_size", default=84, type=int)
    # parser.add_argument("--action_repeat", default=1, type=int)
    parser.add_argument("--frame_stack", default=3, type=int)
    parser.add_argument("--render", action="store_true")
    # replay buffer
    parser.add_argument("--replay_buffer_capacity", default=50000, type=int)
    # train
    parser.add_argument("--agent", default="curl_sac", type=str)
    parser.add_argument("--init_steps", default=10000, type=int)
    parser.add_argument("--num_train_steps", default=3000000, type=int)
    parser.add_argument("--batch_size", default=32, type=int)
    parser.add_argument("--hidden_dim", default=1024, type=int)
    # eval
    parser.add_argument("--eval_freq", default=2000, type=int)
    parser.add_argument("--num_eval_episodes", default=4, type=int)
    # critic
    parser.add_argument("--critic_lr", default=1e-3, type=float)
    parser.add_argument("--critic_beta", default=0.9, type=float)
    parser.add_argument("--critic_tau", default=0.01, type=float)  # try 0.05 or 0.1
    parser.add_argument(
        "--critic_target_update_freq", default=2, type=int
    )  # try to change it to 1 and retain 0.01 above
    # actor
    parser.add_argument("--actor_lr", default=1e-3, type=float)
    parser.add_argument("--actor_beta", default=0.9, type=float)
    parser.add_argument("--actor_log_std_min", default=-10, type=float)
    parser.add_argument("--actor_log_std_max", default=2, type=float)
    parser.add_argument("--actor_update_freq", default=2, type=int)
    # encoder
    parser.add_argument("--encoder_type", default="mixed", type=str)
    parser.add_argument("--encoder_feature_dim", default=50, type=int)
    parser.add_argument("--encoder_lr", default=1e-3, type=float)
    parser.add_argument("--encoder_tau", default=0.05, type=float)
    parser.add_argument("--num_layers", default=4, type=int)
    parser.add_argument("--num_filters", default=32, type=int)
    parser.add_argument("--curl_latent_dim", default=128, type=int)
    # sac
    parser.add_argument("--discount", default=0.99, type=float)
    parser.add_argument("--init_temperature", default=0.1, type=float)
    parser.add_argument("--alpha_lr", default=1e-4, type=float)
    parser.add_argument("--alpha_beta", default=0.5, type=float)
    # misc
    parser.add_argument("--seed", default=-1, type=int)
    parser.add_argument("--work_dir", default="./runs", type=str)
    parser.add_argument("--save_tb", default=False, action="store_true")
    parser.add_argument("--save_buffer", default=False, action="store_true")
    parser.add_argument("--save_video", default=False, action="store_true")
    parser.add_argument("--save_model", default=False, action="store_true")
    parser.add_argument("--detach_encoder", default=False, action="store_true")
    parser.add_argument("--load", type=str)

    parser.add_argument("--log_interval", default=100, type=int)
    args = parser.parse_args()
    return args

keys_to_monitor=[
    'env_reward/apple_pick/tree/min_fruit_dist_reward',
    'env_reward/apple_pick/tree/gripping_fruit_reward',
    # 'env_reward/apple_pick/tree/force_tree_reward',
    # 'env_reward/apple_pick/tree/force_fruit_reward',
    'env_obs/apple_pick/tree/picks'
                ]

def evaluate(env, agent, video, num_episodes, L, step, args):
    all_ep_rewards = []

    def run_eval_loop(sample_stochastically=True):
        start_time = time.time()
        prefix = "stochastic_" if sample_stochastically else ""
        for i in tqdm(range(num_episodes), desc='eval', unit='ep'):
            obs = env.reset()
            video.init(enabled=(i == 0))
            done = False
            episode_reward = 0
            episode_info = defaultdict(int)
            while not done:
                # center crop image
                if args.encoder_type == "mixed":
                    state, img = utils.split_obs(obs)
                    img = utils.center_crop_image(img, args.image_size)
                    obs = utils.combine_obs(state, img)
                with utils.eval_mode(agent):
                    if sample_stochastically:
                        action = agent.sample_action(obs)
                    else:
                        action = agent.select_action(obs)
                obs, reward, done, info = env.step(action)

                for k in keys_to_monitor:
                    episode_info[k] += info[k]
                video.record(env, yaw=i)
                episode_reward += reward

            for k in keys_to_monitor:
                L.log("eval/" + prefix + k, np.sum(episode_info[k]), step)
            video.save("%d.mp4" % step)
            L.log("eval/" + prefix + "episode_reward", episode_reward, step)
            all_ep_rewards.append(episode_reward)

        L.log("eval/" + prefix + "eval_time", time.time() - start_time, step)
        mean_ep_reward = np.mean(all_ep_rewards)
        best_ep_reward = np.max(all_ep_rewards)
        L.log("eval/" + prefix + "mean_episode_reward", mean_ep_reward, step)
        L.log("eval/" + prefix + "best_episode_reward", best_ep_reward, step)

    run_eval_loop(sample_stochastically=False)
    L.dump(step)


def make_agent(obs_shape, action_shape, args, device):
    if args.agent == "curl_sac":
        return CurlSacAgent(
            obs_shape=obs_shape,
            action_shape=action_shape,
            device=device,
            hidden_dim=args.hidden_dim,
            discount=args.discount,
            init_temperature=args.init_temperature,
            alpha_lr=args.alpha_lr,
            alpha_beta=args.alpha_beta,
            actor_lr=args.actor_lr,
            actor_beta=args.actor_beta,
            actor_log_std_min=args.actor_log_std_min,
            actor_log_std_max=args.actor_log_std_max,
            actor_update_freq=args.actor_update_freq,
            critic_lr=args.critic_lr,
            critic_beta=args.critic_beta,
            critic_tau=args.critic_tau,
            critic_target_update_freq=args.critic_target_update_freq,
            encoder_type=args.encoder_type,
            encoder_feature_dim=args.encoder_feature_dim,
            encoder_lr=args.encoder_lr,
            encoder_tau=args.encoder_tau,
            num_layers=args.num_layers,
            num_filters=args.num_filters,
            log_interval=args.log_interval,
            detach_encoder=args.detach_encoder,
            curl_latent_dim=args.curl_latent_dim,
        )
    else:
        assert "agent is not supported: %s" % args.agent


def main():
    import logging
    from rich.logging import RichHandler
    logging.basicConfig(level=logging.INFO, handlers=[RichHandler(rich_tracebacks=True, markup=True)])

    args = parse_args()
    if args.seed == -1:
        args.__dict__["seed"] = np.random.randint(1, 1000000)
        print('seed', args.__dict__["seed"])
    print(args)
    utils.set_seed_everywhere(args.seed)

    env = gym.make(args.domain_name, render=args.render)
    print(env)
    # TODO action repeat wrapper?

    env.seed(args.seed)

    # # stack several consecutive frames together
    if args.encoder_type == "mixed":
        from apple_gym.env.wrappers import FrameStack, ImageState, PermuteImages

        env = FrameStack(
            PermuteImages(ImageState(env), keys=["img"]),
            n=args.frame_stack,
            keys=["img"],
        )

    if args.load =='auto':
        load_dirs = Path(args.work_dir).glob('*/model/curl*.pt')
        load_dirs = sorted(set([str(d.parent) for d in load_dirs]))
        print('load_dirs', load_dirs)
        args.load = str(load_dirs[-1])
        print('auto load', load_dirs)

    # make directory
    ts = time.gmtime()
    ts = time.strftime("%m-%d", ts)
    env_name = args.domain_name
    exp_name = (
        env_name
        + "-"
        + ts
        + "-im"
        + str(args.image_size)
        + "-b"
        + str(args.batch_size)
        + "-s"
        + str(args.seed)
        + "-"
        + args.encoder_type
    )
    args.work_dir = args.work_dir + "/" + exp_name
    print('work_dir', args.work_dir)

    utils.make_dir(args.work_dir)
    video_dir = utils.make_dir(os.path.join(args.work_dir, "video"))
    model_dir = utils.make_dir(os.path.join(args.work_dir, "model"))
    buffer_dir = utils.make_dir(os.path.join(args.work_dir, "buffer"))

    video = VideoRecorder(video_dir if args.save_video else None)

    with open(os.path.join(args.work_dir, "args.json"), "w") as f:
        json.dump(vars(args), f, sort_keys=True, indent=4)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f'device {device}')

    # shapes
    action_shape = env.action_space.shape
    img = env.observation_space.sample()["img"]
    img_aug = utils.center_crop_image(img, args.image_size)
    obs_shape = {"img": img_aug.shape, "state": env.observation_space["state"].shape}

    replay_buffer = utils.ReplayBuffer(
        obs_space=env.observation_space,
        action_space=env.action_space,
        capacity=args.replay_buffer_capacity,
        batch_size=args.batch_size,
        device=device,
        image_size=args.image_size,
    )

    agent = make_agent(
        obs_shape=obs_shape, action_shape=action_shape, args=args, device=device
    )
    if args.load is not None:
        agent.load_curl(args.load)
    

    # summarize
    obs = env.observation_space.sample()
    state, img = utils.split_obs(obs)
    img_crop = utils.center_crop_image(img, agent.image_size)
    obs_crop = utils.combine_obs(state, img_crop)
    obs_crop['img'] = torch.FloatTensor(obs_crop['img']).to(agent.device).unsqueeze(0)
    obs_crop['state'] = torch.FloatTensor(obs_crop['state']).to(agent.device).unsqueeze(0)
    action = agent.sample_action(obs)
    action = torch.FloatTensor(action).to(agent.device).unsqueeze(0)
    from torchsummaryX import summary
    with torch.no_grad():
        print(agent.critic)
        summary(agent.critic, obs_crop, action)
        print(agent.actor)
        summary(agent.actor, obs_crop)

    L = Logger(args.work_dir, use_tb=args.save_tb)

    episode, episode_reward, done = 0, 0, True
    episode_info = defaultdict(int)
    start_time = time.time()

    for step in tqdm(range(args.num_train_steps), desc="train", unit="step", mininterval=360):
        # evaluate agent periodically

        if (step % args.eval_freq == 0) and (step >= args.eval_freq):
            L.log("eval/episode", episode, step)
            evaluate(env, agent, video, args.num_eval_episodes, L, step, args)
            if args.save_model:
                agent.save_curl(model_dir, step)
            if args.save_buffer:
                replay_buffer.save(buffer_dir)

        if done:
            if step > 0:
                if step % args.log_interval == 0:
                    L.log("train/duration", time.time() - start_time, step)
                    L.dump(step)
                start_time = time.time()
            if step % args.log_interval == 0:
                L.log("train/episode_reward", episode_reward, step)

            for k in keys_to_monitor:
                L.log("train/episode_info" + k, np.sum(episode_info[k]), step)

            obs = env.reset()
            assert env.observation_space.contains(
                obs
            ), f"obs should be in space. ob={obs} space={env.observation_space}"
            done = False
            episode_reward = 0
            episode_info = defaultdict(int)
            episode_step = 0
            episode += 1
            if step % args.log_interval == 0:
                L.log("train/episode", episode, step)

        # sample action for data collection
        if step < args.init_steps:
            action = env.action_space.sample()
        else:
            with utils.eval_mode(agent):
                action = agent.sample_action(obs)
        assert env.action_space.contains(
            action
        ), f"obs should be in space. ob={action} space={env.action_space}"

        if step % 10 ==0:
            # run training update
            if step >= args.init_steps:
                num_updates = 1
                for _ in range(num_updates):
                    agent.update(replay_buffer, L, step)

        next_obs, reward, done, info = env.step(action)

        # allow infinite bootstrap
        done_bool = 0 if episode_step + 1 == env._max_episode_steps else float(done)
        episode_reward += reward
        replay_buffer.add(obs, action, reward, next_obs, done_bool)
        for k in keys_to_monitor:
            episode_info[k] += info[k]

        obs = next_obs
        episode_step += 1


if __name__ == "__main__":
    torch.multiprocessing.set_start_method("spawn")

    main()

# import dmc2gym # environment # parser.add_argument("--action_repeat", default=1, type=int) # replay buffer # train # eval # critic # try 0.05 or 0.1 # try to change it to 1 and retain 0.01 above # actor # encoder # sac # misc # 'env_reward/apple_pick/tree/force_tree_reward', # 'env_reward/apple_pick/tree/force_fruit_reward', # center crop image # TODO action repeat wrapper? # # stack several consecutive frames together # make directory # shapes # summarize # evaluate agent periodically # sample action for data collection # run training update # allow infinite bootstrap