misc

.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

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`

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)

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

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

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)