Initial commit

Author: bwfbowen

.gitignore

@@ -0,0 +1,151 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# Mac
+.DS_Store
+
+# data
+data/
+
+# test notebook
+*.ipynb
+
+# 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/
+cover/
+
+# 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/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .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
+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/
+
+# Cython debug symbols
+cython_debug/
+
+# static files generated from Django application using `collectstatic`
+media
+static

README.md

@@ -0,0 +1,55 @@
+# SLAMuZero
+
+## acme-more-mcts
+
+## Installation
+fork [acme-more-mcts](https://github.com/bwfbowen/acme-more-mcts) and run following command from the main directory(where `setup.py` is located):
+```sh
+pip install .[jax,tf,testing,envs]
+```
+
+
+Habitat is also needed:
+```sh
+pip install habitat-sim
+pip install habitat-api
+```
+
+## Setup
+The project requires datasets in a `data` folder in the following format (same as habitat-api):
+```
+SLAMuZero/
+  data/
+    scene_datasets/
+      gibson/
+        Adrian.glb
+        Adrian.navmesh
+        ...
+    datasets/
+      pointnav/
+        gibson/
+          v1/
+            train/
+            val/
+            ...
+```
+Please download the data using the instructions here: https://github.com/facebookresearch/habitat-api#data
+
+## Getting started
+To run the code:
+```python
+python run_acme.py
+```
+
+## Visualization
+Pass `--print_images 1` to plot trajectory
+```sh
+python run_acme.py --print_images 1
+```
+
+And use `draw.py` to generate `.gif`
+```sh
+python draw.py --pic_dir PATH_TO_DUMP 
+```
+
+<img src="./assets/demo_s4.gif" alt="drawing" width="300"/>

assets/demo_s4.gif

@@ -0,0 +1,61 @@
+from PIL import Image
+import os
+import re 
+
+from absl import app
+from absl import flags
+
+flags.DEFINE_string('pic_dir', "~/tmp/dump/nts_eval_img/episodes/1/1", 'task config.')
+flags.DEFINE_string('output_file', '~/tmp/dump/nav.gif', 'output file')
+flags.DEFINE_integer('frame_duration', 50, 'frame duration')
+flags.DEFINE_string('pic_suffix', '.png', 'pic suffix')
+flags.DEFINE_list('max_size', [100,100], 'max size')
+flags.DEFINE_boolean('optimize', True, 'whether to optimize')
+
+FLAGS = flags.FLAGS
+
+
+def sorting_key(filename):
+    match = re.search(r"(\d+)-(\d+)-Vis-(\d+)\.png$", filename)
+    if match:
+        return int(match.group(3))
+    return 0  # Default if the pattern is not found
+
+
+def resize_image(image, max_size):
+    return image.resize(max_size, Image.LANCZOS)
+
+
+def generate_gif_from_pics(
+        pic_dir: str,
+        output_file: str,
+        frame_duration: int = 500,
+        pic_suffix: str = '.png',
+        max_size: tuple = (100, 100),
+        sort_fn: callable = sorting_key,
+        optimize: bool = True,
+    ):
+    # Sort the file names
+    pic_dir = os.path.expanduser(pic_dir)
+    file_names = sorted([os.path.join(pic_dir, f) for f in os.listdir(pic_dir) if f.endswith(pic_suffix)], key=sort_fn)
+
+    # Create a list to hold the images
+    images = []
+
+    # Load each file, resize if necessary, and append to images list
+    for file_name in file_names:
+        with Image.open(file_name) as img:
+            img_resized = resize_image(img, tuple(map(int, max_size)))
+            images.append(img_resized.copy())  # Copy to ensure the file is not left open
+
+    # Create and save the GIF
+    gif_path = os.path.expanduser(output_file)
+    images[0].save(gif_path, save_all=True, append_images=images[1:], optimize=optimize, duration=frame_duration, loop=0)
+
+
+def main(_):
+    generate_gif_from_pics(FLAGS.pic_dir, FLAGS.output_file, FLAGS.frame_duration, FLAGS.pic_suffix, FLAGS.max_size, optimize=FLAGS.optimize)
+
+
+if __name__ == '__main__':
+    app.run(main)

draw.py

@@ -0,0 +1,55 @@
+import torch
+
+from .habitat import construct_envs, Exploration_Env
+
+
+def make_vec_envs(args):
+    envs = construct_envs(args)
+    envs = VecPyTorch(envs, args.device)
+    return envs
+
+
+# Adapted from https://github.com/ikostrikov/pytorch-a2c-ppo-acktr-gail/blob/master/a2c_ppo_acktr/envs.py#L159
+class VecPyTorch():
+
+    def __init__(self, venv, device):
+        self.venv = venv
+        self.num_envs = venv.num_envs
+        self.observation_space = venv.observation_space
+        self.action_space = venv.action_space
+        self.device = device
+
+    def reset(self):
+        obs, info = self.venv.reset()
+        obs = torch.from_numpy(obs).float().to(self.device)
+        return obs, info
+
+    def step_async(self, actions):
+        actions = actions.cpu().numpy()
+        self.venv.step_async(actions)
+
+    def step_wait(self):
+        obs, reward, done, info = self.venv.step_wait()
+        obs = torch.from_numpy(obs).float().to(self.device)
+        reward = torch.from_numpy(reward).float()
+        return obs, reward, done, info
+
+    def step(self, actions):
+        actions = actions.cpu().numpy()
+        obs, reward, done, info = self.venv.step(actions)
+        obs = torch.from_numpy(obs).float().to(self.device)
+        reward = torch.from_numpy(reward).float()
+        return obs, reward, done, info
+
+    def get_rewards(self, inputs):
+        reward = self.venv.get_rewards(inputs)
+        reward = torch.from_numpy(reward).float()
+        return reward
+
+    def get_short_term_goal(self, inputs):
+        stg = self.venv.get_short_term_goal(inputs)
+        stg = torch.from_numpy(stg).float()
+        return stg
+
+    def close(self):
+        return self.venv.close()