├── .github └── workflows │ └── pythontests.yml ├── .gitignore ├── LICENSE ├── README.md ├── install.sh ├── pyproject.toml ├── pytest.ini ├── requirements.txt ├── setup.py └── uniMASK ├── __init__.py ├── batches.py ├── data ├── __init__.py ├── datasets │ ├── __init__.py │ ├── generate_maze2d.py │ ├── maze2d-medium-v1-h200.pkl │ ├── maze2d-umaze-v1-h49.pkl │ ├── maze2d-umaze-v1-h50.pkl │ ├── maze2d-umaze-v1-h9.pkl │ └── parse_d4rl_datasets.py ├── test_data │ ├── 2000_keyenv16x16_10len.pkl │ ├── 2000_keyenv_10len.pkl │ └── __init__.py └── transformer_runs │ └── __init__.py ├── envs ├── __init__.py ├── base_data.py ├── d4rl │ ├── __init__.py │ ├── d4rl_data.py │ └── maze │ │ ├── __init__.py │ │ └── data.py ├── evaluator.py ├── minigrid │ ├── __init__.py │ ├── agents.py │ ├── data.py │ ├── env.py │ ├── good_heatmaps │ │ ├── 1000_heatmap_normFalse.png │ │ ├── 1000_heatmap_normTrue.png │ │ ├── 500_heatmap_normFalse.png │ │ ├── 500_heatmap_normTrue.png │ │ ├── 50_heatmap_normFalse.png │ │ └── 50_heatmap_normTrue.png │ ├── heatmaps.py │ ├── inference.py │ └── viz.py └── mujoco │ ├── maze_env.py │ └── maze_viz.py ├── scripts ├── Parse wandb Maze Experiments.ipynb ├── __init__.py ├── argparsing.py ├── configs.py ├── maze_ft_final.sh ├── minigrid_ft_repro.sh ├── minigrid_heatmap.sh ├── minigrid_repro.sh └── train.py ├── sequences.py ├── sweeps ├── medium_maze_sweep_DT.yaml ├── medium_maze_sweep_all.yaml └── minigrid.yaml ├── tests └── test_transformer.py ├── trainer.py ├── transformer.py ├── transformer_train.py └── utils.py /.github/workflows/pythontests.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/.github/workflows/pythontests.yml -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/.gitignore -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/LICENSE -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/README.md -------------------------------------------------------------------------------- /install.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/install.sh -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/pyproject.toml -------------------------------------------------------------------------------- /pytest.ini: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/pytest.ini -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/requirements.txt -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/setup.py -------------------------------------------------------------------------------- /uniMASK/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /uniMASK/batches.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/batches.py -------------------------------------------------------------------------------- /uniMASK/data/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/data/__init__.py -------------------------------------------------------------------------------- /uniMASK/data/datasets/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /uniMASK/data/datasets/generate_maze2d.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/data/datasets/generate_maze2d.py -------------------------------------------------------------------------------- /uniMASK/data/datasets/maze2d-medium-v1-h200.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/data/datasets/maze2d-medium-v1-h200.pkl -------------------------------------------------------------------------------- /uniMASK/data/datasets/maze2d-umaze-v1-h49.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/data/datasets/maze2d-umaze-v1-h49.pkl -------------------------------------------------------------------------------- /uniMASK/data/datasets/maze2d-umaze-v1-h50.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/data/datasets/maze2d-umaze-v1-h50.pkl -------------------------------------------------------------------------------- /uniMASK/data/datasets/maze2d-umaze-v1-h9.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/data/datasets/maze2d-umaze-v1-h9.pkl -------------------------------------------------------------------------------- /uniMASK/data/datasets/parse_d4rl_datasets.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/data/datasets/parse_d4rl_datasets.py -------------------------------------------------------------------------------- /uniMASK/data/test_data/2000_keyenv16x16_10len.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/data/test_data/2000_keyenv16x16_10len.pkl -------------------------------------------------------------------------------- /uniMASK/data/test_data/2000_keyenv_10len.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/data/test_data/2000_keyenv_10len.pkl -------------------------------------------------------------------------------- /uniMASK/data/test_data/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /uniMASK/data/transformer_runs/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /uniMASK/envs/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/envs/__init__.py -------------------------------------------------------------------------------- /uniMASK/envs/base_data.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/envs/base_data.py -------------------------------------------------------------------------------- /uniMASK/envs/d4rl/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /uniMASK/envs/d4rl/d4rl_data.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/envs/d4rl/d4rl_data.py -------------------------------------------------------------------------------- /uniMASK/envs/d4rl/maze/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /uniMASK/envs/d4rl/maze/data.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/envs/d4rl/maze/data.py -------------------------------------------------------------------------------- /uniMASK/envs/evaluator.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/envs/evaluator.py -------------------------------------------------------------------------------- /uniMASK/envs/minigrid/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /uniMASK/envs/minigrid/agents.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/envs/minigrid/agents.py -------------------------------------------------------------------------------- /uniMASK/envs/minigrid/data.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/envs/minigrid/data.py -------------------------------------------------------------------------------- /uniMASK/envs/minigrid/env.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/envs/minigrid/env.py -------------------------------------------------------------------------------- /uniMASK/envs/minigrid/good_heatmaps/1000_heatmap_normFalse.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/envs/minigrid/good_heatmaps/1000_heatmap_normFalse.png -------------------------------------------------------------------------------- /uniMASK/envs/minigrid/good_heatmaps/1000_heatmap_normTrue.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/envs/minigrid/good_heatmaps/1000_heatmap_normTrue.png -------------------------------------------------------------------------------- /uniMASK/envs/minigrid/good_heatmaps/500_heatmap_normFalse.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/envs/minigrid/good_heatmaps/500_heatmap_normFalse.png -------------------------------------------------------------------------------- /uniMASK/envs/minigrid/good_heatmaps/500_heatmap_normTrue.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/envs/minigrid/good_heatmaps/500_heatmap_normTrue.png -------------------------------------------------------------------------------- /uniMASK/envs/minigrid/good_heatmaps/50_heatmap_normFalse.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/envs/minigrid/good_heatmaps/50_heatmap_normFalse.png -------------------------------------------------------------------------------- /uniMASK/envs/minigrid/good_heatmaps/50_heatmap_normTrue.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/envs/minigrid/good_heatmaps/50_heatmap_normTrue.png -------------------------------------------------------------------------------- /uniMASK/envs/minigrid/heatmaps.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/envs/minigrid/heatmaps.py -------------------------------------------------------------------------------- /uniMASK/envs/minigrid/inference.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/envs/minigrid/inference.py -------------------------------------------------------------------------------- /uniMASK/envs/minigrid/viz.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/envs/minigrid/viz.py -------------------------------------------------------------------------------- /uniMASK/envs/mujoco/maze_env.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/envs/mujoco/maze_env.py -------------------------------------------------------------------------------- /uniMASK/envs/mujoco/maze_viz.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/envs/mujoco/maze_viz.py -------------------------------------------------------------------------------- /uniMASK/scripts/Parse wandb Maze Experiments.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/scripts/Parse wandb Maze Experiments.ipynb -------------------------------------------------------------------------------- /uniMASK/scripts/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /uniMASK/scripts/argparsing.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/scripts/argparsing.py -------------------------------------------------------------------------------- /uniMASK/scripts/configs.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/scripts/configs.py -------------------------------------------------------------------------------- /uniMASK/scripts/maze_ft_final.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/scripts/maze_ft_final.sh -------------------------------------------------------------------------------- /uniMASK/scripts/minigrid_ft_repro.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/scripts/minigrid_ft_repro.sh -------------------------------------------------------------------------------- /uniMASK/scripts/minigrid_heatmap.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/scripts/minigrid_heatmap.sh -------------------------------------------------------------------------------- /uniMASK/scripts/minigrid_repro.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/scripts/minigrid_repro.sh -------------------------------------------------------------------------------- /uniMASK/scripts/train.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/scripts/train.py -------------------------------------------------------------------------------- /uniMASK/sequences.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/sequences.py -------------------------------------------------------------------------------- /uniMASK/sweeps/medium_maze_sweep_DT.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/sweeps/medium_maze_sweep_DT.yaml -------------------------------------------------------------------------------- /uniMASK/sweeps/medium_maze_sweep_all.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/sweeps/medium_maze_sweep_all.yaml -------------------------------------------------------------------------------- /uniMASK/sweeps/minigrid.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/sweeps/minigrid.yaml -------------------------------------------------------------------------------- /uniMASK/tests/test_transformer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/tests/test_transformer.py -------------------------------------------------------------------------------- /uniMASK/trainer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/trainer.py -------------------------------------------------------------------------------- /uniMASK/transformer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/transformer.py -------------------------------------------------------------------------------- /uniMASK/transformer_train.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/transformer_train.py -------------------------------------------------------------------------------- /uniMASK/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/micahcarroll/uniMASK/HEAD/uniMASK/utils.py --------------------------------------------------------------------------------