diff --git a/.gitignore b/.gitignore index 92468f5..855ad5f 100644 --- a/.gitignore +++ b/.gitignore @@ -5,3 +5,206 @@ __pycache__/ scripts/run_*.sh test* *frame2state* +runs/ + + +# Created by https://www.toptal.com/developers/gitignore/api/code,python,jupyternotebooks,windows,linux +# Edit at https://www.toptal.com/developers/gitignore?templates=code,python,jupyternotebooks,windows,linux + +### Code ### +.vscode/* +!.vscode/settings.json +!.vscode/tasks.json +!.vscode/launch.json +!.vscode/extensions.json +*.code-workspace + +### JupyterNotebooks ### +# gitignore template for Jupyter Notebooks +# website: http://jupyter.org/ + +.ipynb_checkpoints +*/.ipynb_checkpoints/* + +# IPython +profile_default/ +ipython_config.py + +# Remove previous ipynb_checkpoints +# git rm -r .ipynb_checkpoints/ + +### Linux ### +*~ + +# temporary files which can be created if a process still has a handle open of a deleted file +.fuse_hidden* + +# KDE directory preferences +.directory + +# Linux trash folder which might appear on any partition or disk +.Trash-* + +# .nfs files are created when an open file is removed but is still being accessed +.nfs* + +### Python ### +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +pip-wheel-metadata/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +pytestdebug.log + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ +doc/_build/ + +# PyBuilder +target/ + +# Jupyter Notebook + +# IPython + +# pyenv +.python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# # Environments +# .env +# .venv +# # env/ +# venv/ +# ENV/ +# env.bak/ +# venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +### Windows ### +# Windows thumbnail cache files +Thumbs.db +Thumbs.db:encryptable +ehthumbs.db +ehthumbs_vista.db + +# Dump file +*.stackdump + +# Folder config file +[Dd]esktop.ini + +# Recycle Bin used on file shares +$RECYCLE.BIN/ + +# Windows Installer files +*.cab +*.msi +*.msix +*.msm +*.msp + +# Windows shortcuts +*.lnk + +# End of https://www.toptal.com/developers/gitignore/api/code,python,jupyternotebooks,windows,linux diff --git a/README.md b/README.md index b4ee31a..4beb1e8 100644 --- a/README.md +++ b/README.md @@ -3,25 +3,25 @@ Adapt to Apple gym TODO: - apple - - make env that returns seperate obs dict - - make a wrapper that flattens act, and normalises act - - make a wrapper that flattens obs... - - mkae a wrapper that does grconv + - [x] make env that returns seperate obs dict + - [x] make a wrapper that flattens act, and normalises act + - [x] make a wrapper that flattens obs... + - make a wrapper that does grconv - norm - [x] space should be acc, but it's not fix in diy gym - - [ ] a wrapper that will normaction - - [ ] a wrapper to norm dict obs - - then norm based on space? + - [x] a wrapper that will normaction + - [x] a wrapper to norm dict obs + - [x] then norm based on space? - this - - make env handle it - - make replay buffer handle it + - [x] make env handle it + - [x] make replay buffer handle it - we take in flat array which contains [...state, ...img_flat] - we reshape in encoder - no resize img - always use pixels of course, otherwise no curl -- [ ] replay buffer store flat... even tho pixel mode -- [ ] custom encoder? +- [x] replay buffer store fl at... even tho pixel mode +- [x] custom encoder? `./scripts/run.sh` diff --git a/logger.py b/logger.py index 18fc30b..5d0f05e 100644 --- a/logger.py +++ b/logger.py @@ -74,7 +74,7 @@ class MetersGroup(object): return template % (key, value) def _dump_to_console(self, data, prefix): - prefix = colored(prefix, 'yellow' if prefix == 'train' else 'green') + # prefix = colored(prefix, 'yellow' if prefix == 'train' else 'green') pieces = ['{:5}'.format(prefix)] for key, disp_key, ty in self._formating: value = data.get(key, 0) diff --git a/scripts/run.sh b/scripts/run.sh index 86029a2..e52bb69 100755 --- a/scripts/run.sh +++ b/scripts/run.sh @@ -1,5 +1,6 @@ #!/bin/bash -CUDA_VISIBLE_DEVICES=1 python \ +set -e +CUDA_VISIBLE_DEVICES=1 /home/wassname/anaconda/envs/diygym4/bin/python \ -m pdb -c continue \ train.py \ --save_tb diff --git a/train.py b/train.py index 4622234..3b67fb9 100644 --- a/train.py +++ b/train.py @@ -8,9 +8,12 @@ import sys import random import time import json + # import dmc2gym import copy +from tqdm.auto import tqdm +from rich import print import utils from logger import Logger from video import VideoRecorder @@ -20,61 +23,64 @@ 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("--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') + 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=10000, type=int) + parser.add_argument("--replay_buffer_capacity", default=30000, type=int) # train - parser.add_argument('--agent', default='curl_sac', type=str) - parser.add_argument('--init_steps', default=1000, type=int) - parser.add_argument('--num_train_steps', default=1000000, type=int) - parser.add_argument('--batch_size', default=32, type=int) - parser.add_argument('--hidden_dim', default=1024, type=int) + parser.add_argument("--agent", default="curl_sac", type=str) + parser.add_argument("--init_steps", default=1000, type=int) + parser.add_argument("--num_train_steps", default=1000000, 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=1000, type=int) - parser.add_argument('--num_eval_episodes', default=10, type=int) + parser.add_argument("--eval_freq", default=10000, 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 + 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) + 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) + 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) + 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("--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('--log_interval', default=100, type=int) + parser.add_argument("--log_interval", default=100, type=int) args = parser.parse_args() return args @@ -84,7 +90,7 @@ def evaluate(env, agent, video, num_episodes, L, step, args): def run_eval_loop(sample_stochastically=True): start_time = time.time() - prefix = 'stochastic_' if sample_stochastically else '' + prefix = "stochastic_" if sample_stochastically else "" for i in range(num_episodes): obs = env.reset() video.init(enabled=(i == 0)) @@ -92,7 +98,7 @@ def evaluate(env, agent, video, num_episodes, L, step, args): episode_reward = 0 while not done: # center crop image - if args.encoder_type == 'mixed': + 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) @@ -101,26 +107,35 @@ def evaluate(env, agent, video, num_episodes, L, step, args): action = agent.sample_action(obs) else: action = agent.select_action(obs) - obs, reward, done, _ = env.step(action) + obs, reward, done, info = env.step(action) + # 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' + # ] + # for k in keys_to_monitor: + # L.log("eval/" + prefix + "episode_reward", episode_reward, step) video.record(env) episode_reward += reward - video.save('%d.mp4' % step) - L.log('eval/' + prefix + 'episode_reward', episode_reward, 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) + + 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) + 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': + if args.agent == "curl_sac": return CurlSacAgent( obs_shape=obs_shape, action_shape=action_shape, @@ -147,65 +162,69 @@ def make_agent(obs_shape, action_shape, args, device): num_filters=args.num_filters, log_interval=args.log_interval, detach_encoder=args.detach_encoder, - curl_latent_dim=args.curl_latent_dim - + curl_latent_dim=args.curl_latent_dim, ) else: - assert 'agent is not supported: %s' % args.agent + assert "agent is not supported: %s" % args.agent + def main(): args = parse_args() - if args.seed == -1: - args.__dict__["seed"] = np.random.randint(1,1000000) + if args.seed == -1: + args.__dict__["seed"] = np.random.randint(1, 1000000) utils.set_seed_everywhere(args.seed) env = gym.make(args.domain_name, render=args.render) # TODO action repeat wrapper? - # env = dmc2gym.make( - # domain_name=args.domain_name, - # task_name=args.task_name, - # seed=args.seed, - # visualize_reward=False, - # from_mixeds=(args.encoder_type == 'mixed'), - # height=args.pre_transform_image_size, - # width=args.pre_transform_image_size, - # frame_skip=args.action_repeat - # ) - + env.seed(args.seed) # # stack several consecutive frames together - if args.encoder_type == 'mixed': + 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']) - + + env = FrameStack( + PermuteImages(ImageState(env), keys=["img"]), + n=args.frame_stack, + keys=["img"], + ) + # make directory - ts = time.gmtime() - ts = time.strftime("%m-%d", ts) + 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 + 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 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_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: + 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') + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") action_shape = env.action_space.shape - # TODO, I need cropped aug shape now... - # TODO I need to make split obs and combine obs? - img = env.observation_space.sample()['img'] + 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} + obs_shape = {"img": img_aug.shape, "state": env.observation_space["state"].shape} replay_buffer = utils.ReplayBuffer( obs_space=env.observation_space, @@ -217,23 +236,37 @@ def main(): ) agent = make_agent( - obs_shape=obs_shape, - action_shape=action_shape, - args=args, - device=device + obs_shape=obs_shape, action_shape=action_shape, args=args, device=device ) + + + # 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 start_time = time.time() - for step in range(args.num_train_steps): + for step in tqdm(range(args.num_train_steps), desc="train", unit="step"): # evaluate agent periodically - if step % args.eval_freq == 0: - L.log('eval/episode', episode, step) - evaluate(env, agent, video, args.num_eval_episodes, L, step,args) + if step % args.eval_freq == 400: + 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: @@ -242,20 +275,22 @@ def main(): if done: if step > 0: if step % args.log_interval == 0: - L.log('train/duration', time.time() - start_time, step) + 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) + L.log("train/episode_reward", episode_reward, step) obs = env.reset() - assert env.observation_space.contains(obs), f'obs should be in space. ob={obs} space={env.observation_space}' + assert env.observation_space.contains( + obs + ), f"obs should be in space. ob={obs} space={env.observation_space}" done = False episode_reward = 0 episode_step = 0 episode += 1 if step % args.log_interval == 0: - L.log('train/episode', episode, step) + L.log("train/episode", episode, step) # sample action for data collection if step < args.init_steps: @@ -263,20 +298,20 @@ def main(): 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}' + assert env.action_space.contains( + action + ), f"obs should be in space. ob={action} space={env.action_space}" # run training update if step >= args.init_steps: - num_updates = 1 + num_updates = 1 for _ in range(num_updates): agent.update(replay_buffer, L, step) - next_obs, reward, done, _ = env.step(action) + next_obs, reward, done, info = env.step(action) - # allow infinit bootstrap - done_bool = 0 if episode_step + 1 == env._max_episode_steps else float( - done - ) + # 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) @@ -284,7 +319,7 @@ def main(): episode_step += 1 -if __name__ == '__main__': - torch.multiprocessing.set_start_method('spawn') +if __name__ == "__main__": + torch.multiprocessing.set_start_method("spawn") main() diff --git a/utils.py b/utils.py index 31ecb42..2ddb2ac 100644 --- a/utils.py +++ b/utils.py @@ -76,7 +76,7 @@ class ReplayBuffer(Dataset): self.image_size = image_size self.transform = transform # the proprioceptive obs is stored as float32, mixeds obs as uint8 - obs_dtype = np.float32 if len(obs_shape) == 1 else np.uint8 + obs_dtype = np.float16 if len(obs_shape) == 1 else np.uint8 self.obses = np.empty((capacity, *obs_shape), dtype=obs_dtype) self.next_obses = np.empty((capacity, *obs_shape), dtype=obs_dtype) @@ -93,6 +93,7 @@ class ReplayBuffer(Dataset): def add(self, obs, action, reward, next_obs, done): obs = flatten(obs) + next_obs = flatten(next_obs) np.copyto(self.obses[self.idx], obs) np.copyto(self.actions[self.idx], action) np.copyto(self.rewards[self.idx], reward)