{"id":18438962,"url":"https://github.com/mila-iqia/atari-representation-learning","last_synced_at":"2025-07-20T02:37:35.934Z","repository":{"id":40659883,"uuid":"167142473","full_name":"mila-iqia/atari-representation-learning","owner":"mila-iqia","description":"Code for \"Unsupervised State Representation Learning in Atari\" ","archived":false,"fork":false,"pushed_at":"2023-11-02T08:10:23.000Z","size":6432,"stargazers_count":250,"open_issues_count":12,"forks_count":52,"subscribers_count":16,"default_branch":"master","last_synced_at":"2025-05-19T19:07:17.817Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":"https://arxiv.org/abs/1906.08226","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/mila-iqia.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null}},"created_at":"2019-01-23T08:01:46.000Z","updated_at":"2025-05-16T07:26:18.000Z","dependencies_parsed_at":"2024-01-18T12:09:16.563Z","dependency_job_id":"9a9823b2-726b-4528-8c88-690b82c36171","html_url":"https://github.com/mila-iqia/atari-representation-learning","commit_stats":null,"previous_names":[],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/mila-iqia/atari-representation-learning","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/mila-iqia%2Fatari-representation-learning","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/mila-iqia%2Fatari-representation-learning/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/mila-iqia%2Fatari-representation-learning/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/mila-iqia%2Fatari-representation-learning/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/mila-iqia","download_url":"https://codeload.github.com/mila-iqia/atari-representation-learning/tar.gz/refs/heads/master","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/mila-iqia%2Fatari-representation-learning/sbom","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":266058481,"owners_count":23870157,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":[],"created_at":"2024-11-06T06:22:51.215Z","updated_at":"2025-07-20T02:37:35.909Z","avatar_url":"https://github.com/mila-iqia.png","language":"Python","funding_links":[],"categories":["Machine Learning"],"sub_categories":["Reinforcement Learning"],"readme":"# Unsupervised State Representation Learning in Atari\n\nAnkesh Anand*, Evan Racah*, Sherjil Ozair*, Yoshua Bengio, Marc-Alexandre Côté, R Devon Hjelm\n\nThis repo provides code for the benchmark and techniques introduced in the paper [Unsupervised State Representation Learning in Atari](https://arxiv.org/abs/1906.08226)\n\n* [📦 Install ](#install) -- Install relevant dependencies and the project\n* [🔧 Usage ](#usage) -- Learn how to use AtariARI and train representations with Spatio-Temporal DeepInfomax (ST-DIM)\n* [🕹️ RAM Annotations ](atariari/benchmark/ram_annotations.py) -- Mapping of RAM indexes to semantic state variables\n* [📜 Poster](Poster.pdf) -- Summary of the paper in a poster format\n* [🎞️ Slides](Slides.pdf) -- Slides describing the paper\n\n## Install\n### AtariARI Wrapper\nYou can do a minimal install to get just the AtariARI (Atari Annotated RAM Interface) wrapper by doing:\n\n```bash\npip install 'gym[atari]'\npip install git+git://github.com/mila-iqia/atari-representation-learning.git\n```\nThis just requires `gym[atari]` and it gives you the ability to play around with the AtariARI wrapper.\nIf you want to use the code for training representation learning methods and probing them, you will need a full installation:\n\n### Full installation (AtariARI Wrapper + Training \u0026 Probing Code)\n\n```bash\n# PyTorch and scikit learn\nconda install pytorch torchvision -c pytorch\nconda install scikit-learn\n\n# Baselines for Atari preprocessing\n# Tensorflow is a dependency, but you don't need to install the GPU version\nconda install tensorflow\npip install git+git://github.com/openai/baselines\n\n# pytorch-a2c-ppo-acktr for RL utils\npip install git+git://github.com/ankeshanand/pytorch-a2c-ppo-acktr-gail\n\n# Clone and install our package\npip install -r requirements.txt\npip install git+git://github.com/mila-iqia/atari-representation-learning.git\n```\n\n## Usage \n### Atari Annotated RAM Interface (AtariARI): \n\n\u003cp align=\"center\"\u003e\n \u003cimg src=\"https://raw.githubusercontent.com/mila-iqia/atari-representation-learning/master/atariari/benchmark/atariari.png\" width=700\u003e\n\u003c/p\u003e\n\nAtariARI exposes the ground truth labels for different state variables for each observation. We have made AtariARI available as a Gym wrapper, to use it simply wrap an Atari gym env with `AtariARIWrapper`. \n\n```python\nimport gym\nfrom atariari.benchmark.wrapper import AtariARIWrapper\nenv = AtariARIWrapper(gym.make('MsPacmanNoFrameskip-v4'))\nobs = env.reset()\nobs, reward, done, info = env.step(1)\n```\n\nNow, `info` is a dictionary of the form:\n\n```python\n{'ale.lives': 3,\n 'labels': {'enemy_sue_x': 88,\n  'enemy_inky_x': 88,\n  'enemy_pinky_x': 88,\n  'enemy_blinky_x': 88,\n  'enemy_sue_y': 80,\n  'enemy_inky_y': 80,\n  'enemy_pinky_y': 80,\n  'enemy_blinky_y': 50,\n  'player_x': 88,\n  'player_y': 98,\n  'fruit_x': 0,\n  'fruit_y': 0,\n  'ghosts_count': 3,\n  'player_direction': 3,\n  'dots_eaten_count': 0,\n  'player_score': 0,\n  'num_lives': 2}}\n```\n**Note:** In our experiments, we use additional preprocessing for Atari environments mainly following Minh et. al, 2014. See [atariari/benchmark/envs.py](atariari/benchmark/envs.py) for more info! \n\nIf you want the raw RAM annotations (which parts of ram correspond to each state variable), check out [atariari/benchmark/ram_annotations.py](atariari/benchmark/ram_annotations.py)\n\n\n### Probing\n----\n⚠️ **Important** ⚠️: The RAM labels are meant for full-sized Atari observations (210 * 160). Probing results won't be accurate if you downsample the observations.\n\nWe provide an interface for the included probing tasks.\n\nFirst, get episodes for train, val and, test:\n\n```python\nfrom atariari.benchmark.episodes import get_episodes\n\ntr_episodes, val_episodes,\\\ntr_labels, val_labels,\\\ntest_episodes, test_labels = get_episodes(env_name=\"PitfallNoFrameskip-v4\", \n                                     steps=50000, \n                                     collect_mode=\"random_agent\")\n```\n\nThen probe them using ProbeTrainer and your encoder (`my_encoder`):\n\n```python\nfrom atariari.benchmark.probe import ProbeTrainer\n\nprobe_trainer = ProbeTrainer(my_encoder, representation_len=my_encoder.feature_size)\nprobe_trainer.train(tr_episodes, val_episodes,\n                     tr_labels, val_labels,)\nfinal_accuracies, final_f1_scores = probe_trainer.test(test_episodes, test_labels)\n```\n\nTo see how we use ProbeTrainer, check out [scripts/run_probe.py](scripts/run_probe.py)\n\nHere is an example of `my_encoder`:\n\n```python \n# get your encoder\nimport torch.nn as nn\nimport torch\nclass MyEncoder(nn.Module):\n    def __init__(self, input_channels, feature_size):\n        super().__init__()\n        self.feature_size = feature_size\n        self.input_channels = input_channels\n        self.final_conv_size = 64 * 9 * 6\n        self.cnn = nn.Sequential(\n            nn.Conv2d(input_channels, 32, 8, stride=4),\n            nn.ReLU(),\n            nn.Conv2d(32, 64, 4, stride=2),\n            nn.ReLU(),\n            nn.Conv2d(64, 128, 4, stride=2),\n            nn.ReLU(),\n            nn.Conv2d(128, 64, 3, stride=1),\n            nn.ReLU()\n        )\n        self.fc = nn.Linear(self.final_conv_size, self.feature_size)\n\n    def forward(self, inputs):\n        x = self.cnn(inputs)\n        x = x.view(x.size(0), -1)\n        return self.fc(x)\n        \n\nmy_encoder = MyEncoder(input_channels=1,feature_size=256)\n# load in weights\nmy_encoder.load_state_dict(torch.load(open(\"path/to/my/weights.pt\", \"rb\")))\n```\n\n### Spatio-Temporal DeepInfoMax:\n\u003cp align=\"center\"\u003e\n \u003cimg src=\"https://raw.githubusercontent.com/mila-iqia/atari-representation-learning/master/atariari/methods/STDIM.png\" width=700\u003e\n\u003c/p\u003e\n\n\n`src/` contains implementations of several representation learning methods, along with `ST-DIM`. Here's a sample usage: \n\n```bash\npython -m scripts.run_probe --method infonce-stdim --env-name {env_name}\n```\nwhere `env_name` is of the form `{game}NoFrameskip-v4`, such as `PongNoFrameskip-v4`\n\n### Citation\n\n```\n@article{anand2019unsupervised,\n  title={Unsupervised State Representation Learning in Atari},\n  author={Anand, Ankesh and Racah, Evan and Ozair, Sherjil and Bengio, Yoshua and C{\\^o}t{\\'e}, Marc-Alexandre and Hjelm, R Devon},\n  journal={arXiv preprint arXiv:1906.08226},\n  year={2019}\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fmila-iqia%2Fatari-representation-learning","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fmila-iqia%2Fatari-representation-learning","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fmila-iqia%2Fatari-representation-learning/lists"}