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()