Initial commit

.gitignore

@@ -0,0 +1,45 @@
+# keep each top level section alphabetical
+# keep each item within the sections alphabetical
+
+# large files
+*.png
+*.pdf
+*.gif
+
+# mac
+*.DS_Store
+
+# misc
+*core
+*logs*
+*runs*
+*tb_logs*
+*wandb
+*wandb_logs*
+outputs*
+
+# python
+*__pycache__
+*.egg-info
+*.hypothesis
+*.ipynb_checkpoints
+*.mypy_cache
+*.pyc
+*.npy
+*.npz
+
+# vim
+*.swp
+
+
+# exorl
+research/exorl/__pycache__/
+research/exorl/.ipynb_checkpoints/
+research/exorl/exp_local
+research/exorl/output
+research/exorl/nbs
+research/exorl/tmp/
+research/exorl/notebooks/
+research/exorl/slurm/
+research/exorl/data/
+research/exorl/datasets/

.isort.cfg

@@ -0,0 +1,6 @@
+[settings]
+use_parentheses=True
+include_trailing_comma=True
+multi_line_output=3
+ensure_newline_before_comments=True
+line_length=88

.pre-commit-config.yaml

@@ -0,0 +1,26 @@
+repos:
+
+# remove unused python imports
+-   repo: https://github.com/myint/autoflake.git
+    rev: v2.0.1
+    hooks:
+    -   id: autoflake
+        args: ["--in-place", "--remove-all-unused-imports", "--ignore-init-module-imports"]
+
+# sort imports
+-   repo: https://github.com/timothycrosley/isort
+    rev: 5.12.0
+    hooks:
+    -   id: isort
+
+# code format according to black
+-   repo: https://github.com/ambv/black
+    rev: 23.1.0
+    hooks:
+    -   id: black
+
+# cleanup notebooks
+-   repo: https://github.com/kynan/nbstripout
+    rev: 0.6.1
+    hooks:
+    -   id: nbstripout

CODE_OF_CONDUCT.md

@@ -0,0 +1,80 @@
+# Code of Conduct
+
+## Our Pledge
+
+In the interest of fostering an open and welcoming environment, we as
+contributors and maintainers pledge to make participation in our project and
+our community a harassment-free experience for everyone, regardless of age, body
+size, disability, ethnicity, sex characteristics, gender identity and expression,
+level of experience, education, socio-economic status, nationality, personal
+appearance, race, religion, or sexual identity and orientation.
+
+## Our Standards
+
+Examples of behavior that contributes to creating a positive environment
+include:
+
+* Using welcoming and inclusive language
+* Being respectful of differing viewpoints and experiences
+* Gracefully accepting constructive criticism
+* Focusing on what is best for the community
+* Showing empathy towards other community members
+
+Examples of unacceptable behavior by participants include:
+
+* The use of sexualized language or imagery and unwelcome sexual attention or
+  advances
+* Trolling, insulting/derogatory comments, and personal or political attacks
+* Public or private harassment
+* Publishing others' private information, such as a physical or electronic
+  address, without explicit permission
+* Other conduct which could reasonably be considered inappropriate in a
+  professional setting
+
+## Our Responsibilities
+
+Project maintainers are responsible for clarifying the standards of acceptable
+behavior and are expected to take appropriate and fair corrective action in
+response to any instances of unacceptable behavior.
+
+Project maintainers have the right and responsibility to remove, edit, or
+reject comments, commits, code, wiki edits, issues, and other contributions
+that are not aligned to this Code of Conduct, or to ban temporarily or
+permanently any contributor for other behaviors that they deem inappropriate,
+threatening, offensive, or harmful.
+
+## Scope
+
+This Code of Conduct applies within all project spaces, and it also applies when
+an individual is representing the project or its community in public spaces.
+Examples of representing a project or community include using an official
+project e-mail address, posting via an official social media account, or acting
+as an appointed representative at an online or offline event. Representation of
+a project may be further defined and clarified by project maintainers.
+
+This Code of Conduct also applies outside the project spaces when there is a
+reasonable belief that an individual's behavior may have a negative impact on
+the project or its community.
+
+## Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported by contacting the project team at <[email protected]>. All
+complaints will be reviewed and investigated and will result in a response that
+is deemed necessary and appropriate to the circumstances. The project team is
+obligated to maintain confidentiality with regard to the reporter of an incident.
+Further details of specific enforcement policies may be posted separately.
+
+Project maintainers who do not follow or enforce the Code of Conduct in good
+faith may face temporary or permanent repercussions as determined by other
+members of the project's leadership.
+
+## Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 1.4,
+available at https://www.contributor-covenant.org/version/1/4/code-of-conduct.html
+
+[homepage]: https://www.contributor-covenant.org
+
+For answers to common questions about this code of conduct, see
+https://www.contributor-covenant.org/faq

CONTRIBUTING.md

@@ -0,0 +1,31 @@
+# Contributing to mtm
+We want to make contributing to this project as easy and transparent as
+possible.
+
+## Pull Requests
+We actively welcome your pull requests.
+
+1. Fork the repo and create your branch from `main`.
+2. If you've added code that should be tested, add tests.
+3. If you've changed APIs, update the documentation.
+4. Ensure the test suite passes.
+5. Make sure your code lints.
+6. If you haven't already, complete the Contributor License Agreement ("CLA").
+
+## Contributor License Agreement ("CLA")
+In order to accept your pull request, we need you to submit a CLA. You only need
+to do this once to work on any of Facebook's open source projects.
+
+Complete your CLA here: <https://code.facebook.com/cla>
+
+## Issues
+We use GitHub issues to track public bugs. Please ensure your description is
+clear and has sufficient instructions to be able to reproduce the issue.
+
+Facebook has a [bounty program](https://www.facebook.com/whitehat/) for the safe
+disclosure of security bugs. In those cases, please go through the process
+outlined on that page and do not file a public issue.
+
+## License
+By contributing to mtm, you agree that your contributions will be licensed
+under the LICENSE file in the root directory of this source tree.

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Meta Research
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,94 @@
+# Masked Trajectory Model
+
+This is the official code base for the paper `Masked Trajectory Models for Prediction, Representation, and Control`
+
+If you find our work useful, consider citing:
+```
+@misc{wu2023mtm,
+  author = {Wu, Philipp and Majumdar, Arjun and Stone, Kevin and Lin, Yixin and Mordatch, Igor and Abbeel, Pieter and Rajeswaran, Aravind},
+  title = {Masked Trajectory Models for Prediction, Representation, and Control},
+  booktitle = {International Conference on Machine Learning},
+  year = {2023},
+}
+```
+
+## Instructions
+
+### Install python packages from scratch
+If you want to make an env from scratch
+
+Make a new conda env
+```
+conda create -n mtm python=3.10
+conda activate mtm
+```
+
+Install torch with gpu
+https://pytorch.org/get-started/locally/
+
+
+Run these commands to install all dependencies
+```
+conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.3 -c pytorch
+pip install -e .
+```
+
+Optionally install dev packages.
+```
+pip install -r requirements_dev.txt
+```
+
+### Adroit Experiments [Optional]
+If you wish to run Adroit experiments, please install also install `robohive`.
+ * [mjrl](https://github.com/aravindr93/mjrl/tree/pvr_beta_1) - use the `pvr_beta_1` branch
+   * `pip install git+https://github.com/aravindr93/mjrl.git@83d35df95eb64274c5e93bb32a0a4e2f6576638a`
+ * [robohive](https://github.com/vikashplus/robohive/tree/stable) - use the `stable` branch (note, it uses gitsubmodules, follow install instructions exactly)
+   * specifically, we have tested using this commit - `c1557f5572977085f053df63f4e81f4b4e1fb17c`
+* Additionally you must download the adroit datasets and put them in the `~/mtm_data` directory
+
+# Running The MTM code
+Example commands can be found in `train_examples.sh`
+An example notebook is located at `example_train_sinusoid.ipynb` which shows a simple example of how MTM can be used for trajectory prediction on a sinusoid dataset.
+
+The main code is located in the `mtm` folder. Here is how you can run some of the experiments.
+ * Simple sinusoidal test data `python research/mtm/train.py +exp_mtm=sinusoid_cont`
+ * D4RL `python research/mtm/train.py +exp_mtm=d4rl_cont`
+ * Adroit `python research/mtm/train.py +exp_mtm=adroit_cont`
+
+### Configuring MTM
+ * The config file for mtm is located at `research/mtm/config.yaml`
+ * Some key parameters
+   * `traj_length`: The length of trajectory sub-segments
+   * `mask_ratios`: A list of mask ratios that is randomly sampled
+   * `mask_pattterns`: A list of masking patterns that are randomly sampled. See `MaskType` under `research/mtm/masks.py` for supported options.
+   * `mode_weights`: (Only applies for `AUTO_MASK`) A list of weights that samples which mode is to be the "autoregressive" one. For example, if the mode order is, `states`, `returns`, `actions`, and mode_weights = [0.2, 0.1, 0.7], then with 0.7 probability, the action token and all future tokens will be masked out.
+
+# Code Organization
+
+### pre-commit hooks
+
+pre-commits hooks are great. This will automatically do some checking/formatting. To use the pre-commit hooks, run the following:
+```
+pip install pre-commit
+pre-commit install
+```
+
+If you want to make a commit without using the pre-commit hook, you can commit with the -n flag (ie. `git commit -n ...`).
+
+### Datasets
+ * all dataset code is located in the `research/mtm/datasets` folder. All datasets have to do is return a pytorch dataset that outputs a dict (named set of trajectories).
+ * a dataset should follow the `DatasetProtocol` specified in `research/mtm/datasets/base.py`.
+ * each dataset should also have corresponding `get_datasets` function where all the dataset specific construction logic happens. This function can take anything as input (as specified in the corresponding `yaml` config) and output the train and val torch `Dataset`.
+
+### Tokenizers
+ * All tokenizer code is found in the `research/mtm/tokenizers` folder. Each tokenizer should inherit from the `Tokenizer` abstract class, found in `research/mtm/tokenizers/base.py`
+ * `Tokenizers` must define a `create` method, which can handle dataset specific construction logic.
+
+# Acknowledgements
+This research would not be possible without building on top of existing open source code. We would like to acknowledge and thank the following:
+ * [FangchenLiu/MaskDP_public](https://github.com/FangchenLiu/MaskDP_public): Masked Decision Prediction, which this work builds upon
+ * [ikostrikov/jaxrl](https://github.com/ikostrikov/jaxrl): A fast Jax library for RL. We used this environment wrapping and data loading code for all d4rl experiments.
+ * [denisyarats/exorl](https://github.com/denisyarats/exorl): ExORL provides datasets collected with unsupervised RL methods which we use in representation learning experiments
+ * [vikashplus/robohive](https://github.com/brentyi/tyro): Provides the Adroit environment
+ * [aravindr93/mjrl](https://github.com/aravindr93/mjrl): Code for training the policy for generating data on Adroit
+ * [brentyi/tyro](https://github.com/brentyi/tyro): Argument parsing and configuration