Code release.

2026-09-09 11:24:31 +08:00 · 2022-09-01 18:02:19 +02:00
commit 75e5ab57e8
56 changed files with 4969 additions and 0 deletions
@@ -0,0 +1,136 @@
+# Custom
+wandb
+.vscode
+outputs
+results/figures
+slurm*
+experiments
+
+# 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/
+
+# 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
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# 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
+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/
@@ -0,0 +1,674 @@
+                    GNU GENERAL PUBLIC LICENSE
+                       Version 3, 29 June 2007
+
+ Copyright (C) 2007 Free Software Foundation, Inc. <https://fsf.org/>
+ Everyone is permitted to copy and distribute verbatim copies
+ of this license document, but changing it is not allowed.
+
+                            Preamble
+
+  The GNU General Public License is a free, copyleft license for
+software and other kinds of works.
+
+  The licenses for most software and other practical works are designed
+to take away your freedom to share and change the works.  By contrast,
+the GNU General Public License is intended to guarantee your freedom to
+share and change all versions of a program--to make sure it remains free
+software for all its users.  We, the Free Software Foundation, use the
+GNU General Public License for most of our software; it applies also to
+any other work released this way by its authors.  You can apply it to
+your programs, too.
+
+  When we speak of free software, we are referring to freedom, not
+price.  Our General Public Licenses are designed to make sure that you
+have the freedom to distribute copies of free software (and charge for
+them if you wish), that you receive source code or can get it if you
+want it, that you can change the software or use pieces of it in new
+free programs, and that you know you can do these things.
+
+  To protect your rights, we need to prevent others from denying you
+these rights or asking you to surrender the rights.  Therefore, you have
+certain responsibilities if you distribute copies of the software, or if
+you modify it: responsibilities to respect the freedom of others.
+
+  For example, if you distribute copies of such a program, whether
+gratis or for a fee, you must pass on to the recipients the same
+freedoms that you received.  You must make sure that they, too, receive
+or can get the source code.  And you must show them these terms so they
+know their rights.
+
+  Developers that use the GNU GPL protect your rights with two steps:
+(1) assert copyright on the software, and (2) offer you this License
+giving you legal permission to copy, distribute and/or modify it.
+
+  For the developers' and authors' protection, the GPL clearly explains
+that there is no warranty for this free software.  For both users' and
+authors' sake, the GPL requires that modified versions be marked as
+changed, so that their problems will not be attributed erroneously to
+authors of previous versions.
+
+  Some devices are designed to deny users access to install or run
+modified versions of the software inside them, although the manufacturer
+can do so.  This is fundamentally incompatible with the aim of
+protecting users' freedom to change the software.  The systematic
+pattern of such abuse occurs in the area of products for individuals to
+use, which is precisely where it is most unacceptable.  Therefore, we
+have designed this version of the GPL to prohibit the practice for those
+products.  If such problems arise substantially in other domains, we
+stand ready to extend this provision to those domains in future versions
+of the GPL, as needed to protect the freedom of users.
+
+  Finally, every program is threatened constantly by software patents.
+States should not allow patents to restrict development and use of
+software on general-purpose computers, but in those that do, we wish to
+avoid the special danger that patents applied to a free program could
+make it effectively proprietary.  To prevent this, the GPL assures that
+patents cannot be used to render the program non-free.
+
+  The precise terms and conditions for copying, distribution and
+modification follow.
+
+                       TERMS AND CONDITIONS
+
+  0. Definitions.
+
+  "This License" refers to version 3 of the GNU General Public License.
+
+  "Copyright" also means copyright-like laws that apply to other kinds of
+works, such as semiconductor masks.
+
+  "The Program" refers to any copyrightable work licensed under this
+License.  Each licensee is addressed as "you".  "Licensees" and
+"recipients" may be individuals or organizations.
+
+  To "modify" a work means to copy from or adapt all or part of the work
+in a fashion requiring copyright permission, other than the making of an
+exact copy.  The resulting work is called a "modified version" of the
+earlier work or a work "based on" the earlier work.
+
+  A "covered work" means either the unmodified Program or a work based
+on the Program.
+
+  To "propagate" a work means to do anything with it that, without
+permission, would make you directly or secondarily liable for
+infringement under applicable copyright law, except executing it on a
+computer or modifying a private copy.  Propagation includes copying,
+distribution (with or without modification), making available to the
+public, and in some countries other activities as well.
+
+  To "convey" a work means any kind of propagation that enables other
+parties to make or receive copies.  Mere interaction with a user through
+a computer network, with no transfer of a copy, is not conveying.
+
+  An interactive user interface displays "Appropriate Legal Notices"
+to the extent that it includes a convenient and prominently visible
+feature that (1) displays an appropriate copyright notice, and (2)
+tells the user that there is no warranty for the work (except to the
+extent that warranties are provided), that licensees may convey the
+work under this License, and how to view a copy of this License.  If
+the interface presents a list of user commands or options, such as a
+menu, a prominent item in the list meets this criterion.
+
+  1. Source Code.
+
+  The "source code" for a work means the preferred form of the work
+for making modifications to it.  "Object code" means any non-source
+form of a work.
+
+  A "Standard Interface" means an interface that either is an official
+standard defined by a recognized standards body, or, in the case of
+interfaces specified for a particular programming language, one that
+is widely used among developers working in that language.
+
+  The "System Libraries" of an executable work include anything, other
+than the work as a whole, that (a) is included in the normal form of
+packaging a Major Component, but which is not part of that Major
+Component, and (b) serves only to enable use of the work with that
+Major Component, or to implement a Standard Interface for which an
+implementation is available to the public in source code form.  A
+"Major Component", in this context, means a major essential component
+(kernel, window system, and so on) of the specific operating system
+(if any) on which the executable work runs, or a compiler used to
+produce the work, or an object code interpreter used to run it.
+
+  The "Corresponding Source" for a work in object code form means all
+the source code needed to generate, install, and (for an executable
+work) run the object code and to modify the work, including scripts to
+control those activities.  However, it does not include the work's
+System Libraries, or general-purpose tools or generally available free
+programs which are used unmodified in performing those activities but
+which are not part of the work.  For example, Corresponding Source
+includes interface definition files associated with source files for
+the work, and the source code for shared libraries and dynamically
+linked subprograms that the work is specifically designed to require,
+such as by intimate data communication or control flow between those
+subprograms and other parts of the work.
+
+  The Corresponding Source need not include anything that users
+can regenerate automatically from other parts of the Corresponding
+Source.
+
+  The Corresponding Source for a work in source code form is that
+same work.
+
+  2. Basic Permissions.
+
+  All rights granted under this License are granted for the term of
+copyright on the Program, and are irrevocable provided the stated
+conditions are met.  This License explicitly affirms your unlimited
+permission to run the unmodified Program.  The output from running a
+covered work is covered by this License only if the output, given its
+content, constitutes a covered work.  This License acknowledges your
+rights of fair use or other equivalent, as provided by copyright law.
+
+  You may make, run and propagate covered works that you do not
+convey, without conditions so long as your license otherwise remains
+in force.  You may convey covered works to others for the sole purpose
+of having them make modifications exclusively for you, or provide you
+with facilities for running those works, provided that you comply with
+the terms of this License in conveying all material for which you do
+not control copyright.  Those thus making or running the covered works
+for you must do so exclusively on your behalf, under your direction
+and control, on terms that prohibit them from making any copies of
+your copyrighted material outside their relationship with you.
+
+  Conveying under any other circumstances is permitted solely under
+the conditions stated below.  Sublicensing is not allowed; section 10
+makes it unnecessary.
+
+  3. Protecting Users' Legal Rights From Anti-Circumvention Law.
+
+  No covered work shall be deemed part of an effective technological
+measure under any applicable law fulfilling obligations under article
+11 of the WIPO copyright treaty adopted on 20 December 1996, or
+similar laws prohibiting or restricting circumvention of such
+measures.
+
+  When you convey a covered work, you waive any legal power to forbid
+circumvention of technological measures to the extent such circumvention
+is effected by exercising rights under this License with respect to
+the covered work, and you disclaim any intention to limit operation or
+modification of the work as a means of enforcing, against the work's
+users, your or third parties' legal rights to forbid circumvention of
+technological measures.
+
+  4. Conveying Verbatim Copies.
+
+  You may convey verbatim copies of the Program's source code as you
+receive it, in any medium, provided that you conspicuously and
+appropriately publish on each copy an appropriate copyright notice;
+keep intact all notices stating that this License and any
+non-permissive terms added in accord with section 7 apply to the code;
+keep intact all notices of the absence of any warranty; and give all
+recipients a copy of this License along with the Program.
+
+  You may charge any price or no price for each copy that you convey,
+and you may offer support or warranty protection for a fee.
+
+  5. Conveying Modified Source Versions.
+
+  You may convey a work based on the Program, or the modifications to
+produce it from the Program, in the form of source code under the
+terms of section 4, provided that you also meet all of these conditions:
+
+    a) The work must carry prominent notices stating that you modified
+    it, and giving a relevant date.
+
+    b) The work must carry prominent notices stating that it is
+    released under this License and any conditions added under section
+    7.  This requirement modifies the requirement in section 4 to
+    "keep intact all notices".
+
+    c) You must license the entire work, as a whole, under this
+    License to anyone who comes into possession of a copy.  This
+    License will therefore apply, along with any applicable section 7
+    additional terms, to the whole of the work, and all its parts,
+    regardless of how they are packaged.  This License gives no
+    permission to license the work in any other way, but it does not
+    invalidate such permission if you have separately received it.
+
+    d) If the work has interactive user interfaces, each must display
+    Appropriate Legal Notices; however, if the Program has interactive
+    interfaces that do not display Appropriate Legal Notices, your
+    work need not make them do so.
+
+  A compilation of a covered work with other separate and independent
+works, which are not by their nature extensions of the covered work,
+and which are not combined with it such as to form a larger program,
+in or on a volume of a storage or distribution medium, is called an
+"aggregate" if the compilation and its resulting copyright are not
+used to limit the access or legal rights of the compilation's users
+beyond what the individual works permit.  Inclusion of a covered work
+in an aggregate does not cause this License to apply to the other
+parts of the aggregate.
+
+  6. Conveying Non-Source Forms.
+
+  You may convey a covered work in object code form under the terms
+of sections 4 and 5, provided that you also convey the
+machine-readable Corresponding Source under the terms of this License,
+in one of these ways:
+
+    a) Convey the object code in, or embodied in, a physical product
+    (including a physical distribution medium), accompanied by the
+    Corresponding Source fixed on a durable physical medium
+    customarily used for software interchange.
+
+    b) Convey the object code in, or embodied in, a physical product
+    (including a physical distribution medium), accompanied by a
+    written offer, valid for at least three years and valid for as
+    long as you offer spare parts or customer support for that product
+    model, to give anyone who possesses the object code either (1) a
+    copy of the Corresponding Source for all the software in the
+    product that is covered by this License, on a durable physical
+    medium customarily used for software interchange, for a price no
+    more than your reasonable cost of physically performing this
+    conveying of source, or (2) access to copy the
+    Corresponding Source from a network server at no charge.
+
+    c) Convey individual copies of the object code with a copy of the
+    written offer to provide the Corresponding Source.  This
+    alternative is allowed only occasionally and noncommercially, and
+    only if you received the object code with such an offer, in accord
+    with subsection 6b.
+
+    d) Convey the object code by offering access from a designated
+    place (gratis or for a charge), and offer equivalent access to the
+    Corresponding Source in the same way through the same place at no
+    further charge.  You need not require recipients to copy the
+    Corresponding Source along with the object code.  If the place to
+    copy the object code is a network server, the Corresponding Source
+    may be on a different server (operated by you or a third party)
+    that supports equivalent copying facilities, provided you maintain
+    clear directions next to the object code saying where to find the
+    Corresponding Source.  Regardless of what server hosts the
+    Corresponding Source, you remain obligated to ensure that it is
+    available for as long as needed to satisfy these requirements.
+
+    e) Convey the object code using peer-to-peer transmission, provided
+    you inform other peers where the object code and Corresponding
+    Source of the work are being offered to the general public at no
+    charge under subsection 6d.
+
+  A separable portion of the object code, whose source code is excluded
+from the Corresponding Source as a System Library, need not be
+included in conveying the object code work.
+
+  A "User Product" is either (1) a "consumer product", which means any
+tangible personal property which is normally used for personal, family,
+or household purposes, or (2) anything designed or sold for incorporation
+into a dwelling.  In determining whether a product is a consumer product,
+doubtful cases shall be resolved in favor of coverage.  For a particular
+product received by a particular user, "normally used" refers to a
+typical or common use of that class of product, regardless of the status
+of the particular user or of the way in which the particular user
+actually uses, or expects or is expected to use, the product.  A product
+is a consumer product regardless of whether the product has substantial
+commercial, industrial or non-consumer uses, unless such uses represent
+the only significant mode of use of the product.
+
+  "Installation Information" for a User Product means any methods,
+procedures, authorization keys, or other information required to install
+and execute modified versions of a covered work in that User Product from
+a modified version of its Corresponding Source.  The information must
+suffice to ensure that the continued functioning of the modified object
+code is in no case prevented or interfered with solely because
+modification has been made.
+
+  If you convey an object code work under this section in, or with, or
+specifically for use in, a User Product, and the conveying occurs as
+part of a transaction in which the right of possession and use of the
+User Product is transferred to the recipient in perpetuity or for a
+fixed term (regardless of how the transaction is characterized), the
+Corresponding Source conveyed under this section must be accompanied
+by the Installation Information.  But this requirement does not apply
+if neither you nor any third party retains the ability to install
+modified object code on the User Product (for example, the work has
+been installed in ROM).
+
+  The requirement to provide Installation Information does not include a
+requirement to continue to provide support service, warranty, or updates
+for a work that has been modified or installed by the recipient, or for
+the User Product in which it has been modified or installed.  Access to a
+network may be denied when the modification itself materially and
+adversely affects the operation of the network or violates the rules and
+protocols for communication across the network.
+
+  Corresponding Source conveyed, and Installation Information provided,
+in accord with this section must be in a format that is publicly
+documented (and with an implementation available to the public in
+source code form), and must require no special password or key for
+unpacking, reading or copying.
+
+  7. Additional Terms.
+
+  "Additional permissions" are terms that supplement the terms of this
+License by making exceptions from one or more of its conditions.
+Additional permissions that are applicable to the entire Program shall
+be treated as though they were included in this License, to the extent
+that they are valid under applicable law.  If additional permissions
+apply only to part of the Program, that part may be used separately
+under those permissions, but the entire Program remains governed by
+this License without regard to the additional permissions.
+
+  When you convey a copy of a covered work, you may at your option
+remove any additional permissions from that copy, or from any part of
+it.  (Additional permissions may be written to require their own
+removal in certain cases when you modify the work.)  You may place
+additional permissions on material, added by you to a covered work,
+for which you have or can give appropriate copyright permission.
+
+  Notwithstanding any other provision of this License, for material you
+add to a covered work, you may (if authorized by the copyright holders of
+that material) supplement the terms of this License with terms:
+
+    a) Disclaiming warranty or limiting liability differently from the
+    terms of sections 15 and 16 of this License; or
+
+    b) Requiring preservation of specified reasonable legal notices or
+    author attributions in that material or in the Appropriate Legal
+    Notices displayed by works containing it; or
+
+    c) Prohibiting misrepresentation of the origin of that material, or
+    requiring that modified versions of such material be marked in
+    reasonable ways as different from the original version; or
+
+    d) Limiting the use for publicity purposes of names of licensors or
+    authors of the material; or
+
+    e) Declining to grant rights under trademark law for use of some
+    trade names, trademarks, or service marks; or
+
+    f) Requiring indemnification of licensors and authors of that
+    material by anyone who conveys the material (or modified versions of
+    it) with contractual assumptions of liability to the recipient, for
+    any liability that these contractual assumptions directly impose on
+    those licensors and authors.
+
+  All other non-permissive additional terms are considered "further
+restrictions" within the meaning of section 10.  If the Program as you
+received it, or any part of it, contains a notice stating that it is
+governed by this License along with a term that is a further
+restriction, you may remove that term.  If a license document contains
+a further restriction but permits relicensing or conveying under this
+License, you may add to a covered work material governed by the terms
+of that license document, provided that the further restriction does
+not survive such relicensing or conveying.
+
+  If you add terms to a covered work in accord with this section, you
+must place, in the relevant source files, a statement of the
+additional terms that apply to those files, or a notice indicating
+where to find the applicable terms.
+
+  Additional terms, permissive or non-permissive, may be stated in the
+form of a separately written license, or stated as exceptions;
+the above requirements apply either way.
+
+  8. Termination.
+
+  You may not propagate or modify a covered work except as expressly
+provided under this License.  Any attempt otherwise to propagate or
+modify it is void, and will automatically terminate your rights under
+this License (including any patent licenses granted under the third
+paragraph of section 11).
+
+  However, if you cease all violation of this License, then your
+license from a particular copyright holder is reinstated (a)
+provisionally, unless and until the copyright holder explicitly and
+finally terminates your license, and (b) permanently, if the copyright
+holder fails to notify you of the violation by some reasonable means
+prior to 60 days after the cessation.
+
+  Moreover, your license from a particular copyright holder is
+reinstated permanently if the copyright holder notifies you of the
+violation by some reasonable means, this is the first time you have
+received notice of violation of this License (for any work) from that
+copyright holder, and you cure the violation prior to 30 days after
+your receipt of the notice.
+
+  Termination of your rights under this section does not terminate the
+licenses of parties who have received copies or rights from you under
+this License.  If your rights have been terminated and not permanently
+reinstated, you do not qualify to receive new licenses for the same
+material under section 10.
+
+  9. Acceptance Not Required for Having Copies.
+
+  You are not required to accept this License in order to receive or
+run a copy of the Program.  Ancillary propagation of a covered work
+occurring solely as a consequence of using peer-to-peer transmission
+to receive a copy likewise does not require acceptance.  However,
+nothing other than this License grants you permission to propagate or
+modify any covered work.  These actions infringe copyright if you do
+not accept this License.  Therefore, by modifying or propagating a
+covered work, you indicate your acceptance of this License to do so.
+
+  10. Automatic Licensing of Downstream Recipients.
+
+  Each time you convey a covered work, the recipient automatically
+receives a license from the original licensors, to run, modify and
+propagate that work, subject to this License.  You are not responsible
+for enforcing compliance by third parties with this License.
+
+  An "entity transaction" is a transaction transferring control of an
+organization, or substantially all assets of one, or subdividing an
+organization, or merging organizations.  If propagation of a covered
+work results from an entity transaction, each party to that
+transaction who receives a copy of the work also receives whatever
+licenses to the work the party's predecessor in interest had or could
+give under the previous paragraph, plus a right to possession of the
+Corresponding Source of the work from the predecessor in interest, if
+the predecessor has it or can get it with reasonable efforts.
+
+  You may not impose any further restrictions on the exercise of the
+rights granted or affirmed under this License.  For example, you may
+not impose a license fee, royalty, or other charge for exercise of
+rights granted under this License, and you may not initiate litigation
+(including a cross-claim or counterclaim in a lawsuit) alleging that
+any patent claim is infringed by making, using, selling, offering for
+sale, or importing the Program or any portion of it.
+
+  11. Patents.
+
+  A "contributor" is a copyright holder who authorizes use under this
+License of the Program or a work on which the Program is based.  The
+work thus licensed is called the contributor's "contributor version".
+
+  A contributor's "essential patent claims" are all patent claims
+owned or controlled by the contributor, whether already acquired or
+hereafter acquired, that would be infringed by some manner, permitted
+by this License, of making, using, or selling its contributor version,
+but do not include claims that would be infringed only as a
+consequence of further modification of the contributor version.  For
+purposes of this definition, "control" includes the right to grant
+patent sublicenses in a manner consistent with the requirements of
+this License.
+
+  Each contributor grants you a non-exclusive, worldwide, royalty-free
+patent license under the contributor's essential patent claims, to
+make, use, sell, offer for sale, import and otherwise run, modify and
+propagate the contents of its contributor version.
+
+  In the following three paragraphs, a "patent license" is any express
+agreement or commitment, however denominated, not to enforce a patent
+(such as an express permission to practice a patent or covenant not to
+sue for patent infringement).  To "grant" such a patent license to a
+party means to make such an agreement or commitment not to enforce a
+patent against the party.
+
+  If you convey a covered work, knowingly relying on a patent license,
+and the Corresponding Source of the work is not available for anyone
+to copy, free of charge and under the terms of this License, through a
+publicly available network server or other readily accessible means,
+then you must either (1) cause the Corresponding Source to be so
+available, or (2) arrange to deprive yourself of the benefit of the
+patent license for this particular work, or (3) arrange, in a manner
+consistent with the requirements of this License, to extend the patent
+license to downstream recipients.  "Knowingly relying" means you have
+actual knowledge that, but for the patent license, your conveying the
+covered work in a country, or your recipient's use of the covered work
+in a country, would infringe one or more identifiable patents in that
+country that you have reason to believe are valid.
+
+  If, pursuant to or in connection with a single transaction or
+arrangement, you convey, or propagate by procuring conveyance of, a
+covered work, and grant a patent license to some of the parties
+receiving the covered work authorizing them to use, propagate, modify
+or convey a specific copy of the covered work, then the patent license
+you grant is automatically extended to all recipients of the covered
+work and works based on it.
+
+  A patent license is "discriminatory" if it does not include within
+the scope of its coverage, prohibits the exercise of, or is
+conditioned on the non-exercise of one or more of the rights that are
+specifically granted under this License.  You may not convey a covered
+work if you are a party to an arrangement with a third party that is
+in the business of distributing software, under which you make payment
+to the third party based on the extent of your activity of conveying
+the work, and under which the third party grants, to any of the
+parties who would receive the covered work from you, a discriminatory
+patent license (a) in connection with copies of the covered work
+conveyed by you (or copies made from those copies), or (b) primarily
+for and in connection with specific products or compilations that
+contain the covered work, unless you entered into that arrangement,
+or that patent license was granted, prior to 28 March 2007.
+
+  Nothing in this License shall be construed as excluding or limiting
+any implied license or other defenses to infringement that may
+otherwise be available to you under applicable patent law.
+
+  12. No Surrender of Others' Freedom.
+
+  If conditions are imposed on you (whether by court order, agreement or
+otherwise) that contradict the conditions of this License, they do not
+excuse you from the conditions of this License.  If you cannot convey a
+covered work so as to satisfy simultaneously your obligations under this
+License and any other pertinent obligations, then as a consequence you may
+not convey it at all.  For example, if you agree to terms that obligate you
+to collect a royalty for further conveying from those to whom you convey
+the Program, the only way you could satisfy both those terms and this
+License would be to refrain entirely from conveying the Program.
+
+  13. Use with the GNU Affero General Public License.
+
+  Notwithstanding any other provision of this License, you have
+permission to link or combine any covered work with a work licensed
+under version 3 of the GNU Affero General Public License into a single
+combined work, and to convey the resulting work.  The terms of this
+License will continue to apply to the part which is the covered work,
+but the special requirements of the GNU Affero General Public License,
+section 13, concerning interaction through a network will apply to the
+combination as such.
+
+  14. Revised Versions of this License.
+
+  The Free Software Foundation may publish revised and/or new versions of
+the GNU General Public License from time to time.  Such new versions will
+be similar in spirit to the present version, but may differ in detail to
+address new problems or concerns.
+
+  Each version is given a distinguishing version number.  If the
+Program specifies that a certain numbered version of the GNU General
+Public License "or any later version" applies to it, you have the
+option of following the terms and conditions either of that numbered
+version or of any later version published by the Free Software
+Foundation.  If the Program does not specify a version number of the
+GNU General Public License, you may choose any version ever published
+by the Free Software Foundation.
+
+  If the Program specifies that a proxy can decide which future
+versions of the GNU General Public License can be used, that proxy's
+public statement of acceptance of a version permanently authorizes you
+to choose that version for the Program.
+
+  Later license versions may give you additional or different
+permissions.  However, no additional obligations are imposed on any
+author or copyright holder as a result of your choosing to follow a
+later version.
+
+  15. Disclaimer of Warranty.
+
+  THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY
+APPLICABLE LAW.  EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT
+HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY
+OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO,
+THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+PURPOSE.  THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM
+IS WITH YOU.  SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF
+ALL NECESSARY SERVICING, REPAIR OR CORRECTION.
+
+  16. Limitation of Liability.
+
+  IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING
+WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS
+THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY
+GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE
+USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF
+DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD
+PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS),
+EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF
+SUCH DAMAGES.
+
+  17. Interpretation of Sections 15 and 16.
+
+  If the disclaimer of warranty and limitation of liability provided
+above cannot be given local legal effect according to their terms,
+reviewing courts shall apply local law that most closely approximates
+an absolute waiver of all civil liability in connection with the
+Program, unless a warranty or assumption of liability accompanies a
+copy of the Program in return for a fee.
+
+                     END OF TERMS AND CONDITIONS
+
+            How to Apply These Terms to Your New Programs
+
+  If you develop a new program, and you want it to be of the greatest
+possible use to the public, the best way to achieve this is to make it
+free software which everyone can redistribute and change under these terms.
+
+  To do so, attach the following notices to the program.  It is safest
+to attach them to the start of each source file to most effectively
+state the exclusion of warranty; and each file should have at least
+the "copyright" line and a pointer to where the full notice is found.
+
+    <one line to give the program's name and a brief idea of what it does.>
+    Copyright (C) <year>  <name of author>
+
+    This program is free software: you can redistribute it and/or modify
+    it under the terms of the GNU General Public License as published by
+    the Free Software Foundation, either version 3 of the License, or
+    (at your option) any later version.
+
+    This program is distributed in the hope that it will be useful,
+    but WITHOUT ANY WARRANTY; without even the implied warranty of
+    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+    GNU General Public License for more details.
+
+    You should have received a copy of the GNU General Public License
+    along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+Also add information on how to contact you by electronic and paper mail.
+
+  If the program does terminal interaction, make it output a short
+notice like this when it starts in an interactive mode:
+
+    <program>  Copyright (C) <year>  <name of author>
+    This program comes with ABSOLUTELY NO WARRANTY; for details type `show w'.
+    This is free software, and you are welcome to redistribute it
+    under certain conditions; type `show c' for details.
+
+The hypothetical commands `show w' and `show c' should show the appropriate
+parts of the General Public License.  Of course, your program's commands
+might be different; for a GUI interface, you would use an "about box".
+
+  You should also get your employer (if you work as a programmer) or school,
+if any, to sign a "copyright disclaimer" for the program, if necessary.
+For more information on this, and how to apply and follow the GNU GPL, see
+<https://www.gnu.org/licenses/>.
+
+  The GNU General Public License does not permit incorporating your program
+into proprietary programs.  If your program is a subroutine library, you
+may consider it more useful to permit linking proprietary applications with
+the library.  If this is what you want to do, use the GNU Lesser General
+Public License instead of this License.  But first, please read
+<https://www.gnu.org/licenses/why-not-lgpl.html>.
@@ -0,0 +1,113 @@
+# Transformers are Sample Efficient World Models (IRIS)
+
+[Transformers are Sample Efficient World Models](https://arXiv-url.com) <br>
+[Vincent Micheli](https://vmicheli.github.io)\*, [Eloi Alonso](https://eloialonso.github.io)\*, [François Fleuret](https://fleuret.org/francois/) <br>
+\* Denotes equal contribution
+
+
+<div align='center'>
+  IRIS agent after 100k environment steps, i.e. two hours of real-time experience
+  <img alt="IRIS playing on Asterix, Boxing, Breakout, Demon Attack, Freeway, Gopher, Kung Fu Master, Pong" src="assets/iris.gif">
+</div>
+
+**tl;dr**
+
+- IRIS is a data-efficient agent trained over millions of imagined trajectories in a world model.
+- The world model is composed of a discrete autoencoder and an autoregressive Transformer.
+- Our approach casts dynamics learning as a sequence modeling problem, where the autoencoder builds a language of image tokens and the Transformer composes that language over time.
+
+
+
+## BibTeX
+
+If you find this code or paper useful, please use the following reference:
+
+```
+Coming soon
+```
+
+## Setup
+
+- Install [PyTorch](https://pytorch.org/get-started/locally/) (torch and torchvision). Code developed with torch==1.11.0 and torchvision==0.12.0.
+- Install [other dependencies](requirements.txt): `pip install -r requirements.txt`
+- Warning: Atari ROMs will be downloaded with the dependencies, which means that you acknowledge that you have the license to use them.
+
+## Launch a training run
+
+```bash
+python src/main.py env.train.id=BreakoutNoFrameskip-v4 common.device=cuda:0 wandb.mode=online
+```
+
+By default, the logs are synced to [weights & biases](https://wandb.ai), set `wandb.mode=disabled` to turn it off.
+
+## Configuration
+
+- All configuration files are located in `config/`, the main configuration file is `config/trainer.yaml`.
+- The simplest way to customize the configuration is to edit these files directly.
+- Please refer to [Hydra](https://github.com/facebookresearch/hydra) for more details regarding configuration management.
+
+## Run folder
+
+Each new run is located at `outputs/YYYY-MM-DD/hh-mm-ss/`. This folder is structured as:
+
+```txt
+outputs/YYYY-MM-DD/hh-mm-ss/
+│
+└─── checkpoints
+│   │   last.pt
+|   |   optimizer.pt
+|   |   ...
+│   │
+│   └─── dataset
+│       │   0.pt
+│       │   1.pt
+│       │   ...
+│
+└─── config
+│   |   trainer.yaml
+|
+└─── media
+│   │
+│   └─── episodes
+│   |   │   ...
+│   │
+│   └─── reconstructions
+│   |   │   ...
+│
+└─── scripts
+|   |   eval.py
+│   │   play.sh
+│   │   resume.sh
+|   |   ...
+|
+└─── src
+|   |   ...
+|
+└─── wandb
+    |   ...
+```
+
+- `checkpoints`: contains the last checkpoint of the model, its optimizer and the dataset.
+- `media`:
+  - `episodes`: contains train / test / imagination episodes for visualization purposes.
+  - `reconstructions`: contains original frames alongside their reconstructions with the autoencoder.
+- `scripts`: **from the run folder**, you can use the following three scripts.
+  - `eval.py`: Launch `python ./scripts/eval.py` to evaluate the run.
+  - `resume.sh`: Launch `./scripts/resume.sh` to resume a training that crashed.
+  - `play.sh`: Tool to visualize some interesting aspects of the run.
+    - Launch `./scripts/play.sh -a` to watch the agent play live in the environment. The left panel displays the original environment, and the right panel shows what the agent actually sees through its discrete autoencoder.
+    - Launch `./scripts/play.sh -w` to unroll live trajectories with your keyboard inputs (i.e. to play in the world model). Note that for faster interaction, the memory of the Transformer is flushed every 20 frames.
+    - Launch `./scripts/play.sh` to visualize the episodes contained in `media/episodes`.
+
+## Results notebook
+
+The folder `results/data/` contains raw scores (for each game, and for each training run) for IRIS and the baselines.
+
+Use the notebook `results/results_iris.ipynb` to reproduce the figures from the paper.
+
+## Credits
+
+- [https://github.com/pytorch/pytorch](https://github.com/pytorch/pytorch)
+- [https://github.com/CompVis/taming-transformers](https://github.com/CompVis/taming-transformers)
+- [https://github.com/karpathy/minGPT](https://github.com/karpathy/minGPT)
+- [https://github.com/google-research/rliable](https://github.com/google-research/rliable)
@@ -0,0 +1 @@
+use_original_obs: False
@@ -0,0 +1,8 @@
+train:
+  _target_: dataset.EpisodesDatasetRamMonitoring
+  max_ram_usage: 30G
+  name: train_dataset
+test:
+  _target_: dataset.EpisodesDataset
+  max_num_episodes: null
+  name: test_dataset
@@ -0,0 +1,21 @@
+train:
+  _target_: envs.make_atari
+  id: null
+  size: 64
+  max_episode_steps: 20000
+  noop_max: 30
+  frame_skip: 4
+  done_on_life_loss: True
+  clip_reward: False
+
+test:
+  _target_: ${..train._target_}
+  id: ${..train.id}
+  size: ${..train.size}
+  max_episode_steps: 108000
+  noop_max: 1
+  frame_skip: ${..train.frame_skip}
+  done_on_life_loss: False
+  clip_reward: False
+
+keymap: atari/${.train.id}
@@ -0,0 +1,20 @@
+_target_: models.tokenizer.Tokenizer
+
+vocab_size: 512
+embed_dim: 512
+encoder:
+  _target_: models.tokenizer.Encoder
+  config:
+    _target_: models.tokenizer.EncoderDecoderConfig
+    resolution: 64
+    in_channels: 3
+    z_channels: 512
+    ch: 64
+    ch_mult: [1, 1, 1, 1, 1]
+    num_res_blocks: 2
+    attn_resolutions: [8, 16]
+    out_ch: 3
+    dropout: 0.0
+decoder:
+  _target_: models.tokenizer.Decoder
+  config: ${..encoder.config}
@@ -0,0 +1,95 @@
+defaults:
+  - _self_
+  - tokenizer: default
+  - world_model: default
+  - actor_critic: default
+  - env: default
+  - datasets: default
+
+wandb:
+  mode: online
+  project: iris
+  entity: null
+  name: null
+  group: null
+  tags: null
+  notes: null
+
+initialization:
+  path_to_checkpoint: null
+  load_tokenizer: False
+  load_world_model: False
+  load_actor_critic: False
+
+common:
+  epochs: 600
+  device: cuda:0
+  do_checkpoint: True
+  seed: 0
+  sequence_length: ${world_model.max_blocks}
+  resume: False # set by resume.sh script only.
+
+collection:
+  train:
+    num_envs: 1
+    stop_after_epochs: 500
+    num_episodes_to_save: 10
+    config:
+      epsilon: 0.01
+      should_sample: True
+      temperature: 1.0
+      num_steps: 200
+      burn_in: ${training.actor_critic.burn_in}
+  test:
+    num_envs: 8
+    num_episodes_to_save: ${collection.train.num_episodes_to_save}
+    config:
+      epsilon: 0.0
+      should_sample: True
+      temperature: 0.5
+      num_episodes: 16
+      burn_in: ${training.actor_critic.burn_in}
+
+training:
+  should: True
+  learning_rate: 0.0001
+  sampling_weights: [0.125, 0.125, 0.25, 0.5]
+  tokenizer:
+    batch_num_samples: 256
+    grad_acc_steps: 1
+    max_grad_norm: 10.0
+    start_after_epochs: 5
+    steps_per_epoch: 200
+  world_model:
+    batch_num_samples: 64
+    grad_acc_steps: 1
+    max_grad_norm: 10.0
+    weight_decay: 0.01
+    start_after_epochs: 25
+    steps_per_epoch: 200
+  actor_critic:
+    batch_num_samples: 64
+    grad_acc_steps: 1
+    max_grad_norm: 10.0
+    start_after_epochs: 50
+    steps_per_epoch: 200
+    imagine_horizon: ${common.sequence_length}
+    burn_in: 20
+    gamma: 0.995
+    lambda_: 0.95
+    entropy_weight: 0.001
+
+evaluation:
+  should: True
+  every: 5
+  tokenizer:
+    batch_num_samples: ${training.tokenizer.batch_num_samples}
+    start_after_epochs: ${training.tokenizer.start_after_epochs}
+    save_reconstructions: True
+  world_model:
+    batch_num_samples: ${training.world_model.batch_num_samples}
+    start_after_epochs: ${training.world_model.start_after_epochs}
+  actor_critic:
+    num_episodes_to_save: ${training.actor_critic.batch_num_samples}
+    horizon: ${training.actor_critic.imagine_horizon}
+    start_after_epochs: ${training.actor_critic.start_after_epochs}
@@ -0,0 +1,10 @@
+_target_: models.TransformerConfig
+tokens_per_block: 17
+max_blocks: 20
+attention: 'causal'
+num_layers: 10
+num_heads: 4
+embed_dim: 256
+embed_pdrop: 0.1
+resid_pdrop: 0.1
+attn_pdrop: 0.1
@@ -0,0 +1,10 @@
+ale-py==0.7.4
+einops==0.3.2
+gym[accept-rom-license]==0.21.0
+hydra-core==1.1.1
+opencv-python
+protobuf==3.20.*
+psutil==5.8.0
+pygame==2.1.2
+tqdm==4.62.3
+wandb==0.12.7
@@ -0,0 +1 @@
+{"Alien": [808.5], "Amidar": [148.6], "Assault": [1263.1], "Asterix": [25557.8], "BankHeist": [351.0], "BattleZone": [13871.2], "Boxing": [52.7], "Breakout": [414.1], "ChopperCommand": [1117.3], "CrazyClimber": [83940.2], "DemonAttack": [13003.9], "Freeway": [21.8], "Frostbite": [296.3], "Gopher": [3260.3], "Hero": [9315.9], "Jamesbond": [517.0], "Kangaroo": [724.1], "Krull": [5663.3], "KungFuMaster": [30944.8], "MsPacman": [1281.2], "Pong": [20.1], "PrivateEye": [96.7], "Qbert": [13781.9], "RoadRunner": [17751.3], "Seaquest": [1100.2], "UpNDown": [17264.2]}
@@ -0,0 +1 @@
+{"Alien": [7127.7], "Amidar": [1719.5], "Assault": [742.0], "Asterix": [8503.3], "BankHeist": [753.1], "BattleZone": [37187.5], "Boxing": [12.1], "Breakout": [30.5], "ChopperCommand": [7387.8], "CrazyClimber": [35829.4], "DemonAttack": [1971.0], "Freeway": [29.6], "Frostbite": [4334.7], "Gopher": [2412.5], "Hero": [30826.4], "Jamesbond": [302.8], "Kangaroo": [3035.0], "Krull": [2665.5], "KungFuMaster": [22736.3], "MsPacman": [6951.6], "Pong": [14.6], "PrivateEye": [69571.3], "Qbert": [13455.0], "RoadRunner": [7845.0], "Seaquest": [42054.7], "UpNDown": [11693.2]}
@@ -0,0 +1,184 @@
+{
+    "Alien": [
+        368.7,
+        452.5,
+        437.2,
+        331.7,
+        510
+    ],
+    "Amidar": [
+        196.0,
+        126,
+        164.9,
+        162.9,
+        65.1
+    ],
+    "Assault": [
+        1481.3,
+        928.1,
+        1121.8,
+        1556.0,
+        2534.8
+    ],
+    "Asterix": [
+        762.5,
+        1114,
+        1082,
+        716.5,
+        593
+    ],
+    "BankHeist": [
+        42.2,
+        48.8,
+        68.4,
+        46.4,
+        59.6
+    ],
+    "BattleZone": [
+        16990,
+        14570,
+        11300,
+        10070,
+        12440
+    ],
+    "Boxing": [
+        86.2,
+        49.3,
+        55.8,
+        83.5,
+        75.5
+    ],
+    "Breakout": [
+        104.6,
+        85.6,
+        69.7,
+        70.3,
+        88.1
+    ],
+    "ChopperCommand": [
+        1150,
+        1276,
+        1659,
+        2162,
+        1578
+    ],
+    "CrazyClimber": [
+        64101,
+        66127,
+        62740,
+        77066,
+        26587
+    ],
+    "DemonAttack": [
+        905,
+        2011.9,
+        2888.5,
+        2613.4,
+        1753
+    ],
+    "Freeway": [
+        30.9,
+        31.5,
+        30.5,
+        30.8,
+        31.6
+    ],
+    "Frostbite": [
+        274.3,
+        250.0,
+        295,
+        250,
+        226.3
+    ],
+    "Gopher": [
+        2987.2,
+        2476.8,
+        3110.4,
+        1171.4,
+        1434.8
+    ],
+    "Hero": [
+        6916.2,
+        6933.9,
+        6740.6,
+        7881.6,
+        6714.7
+    ],
+    "Jamesbond": [
+        487.5,
+        394,
+        486,
+        464,
+        482
+    ],
+    "Kangaroo": [
+        2249,
+        732,
+        750,
+        90,
+        370
+    ],
+    "Krull": [
+        6047.4,
+        7364.8,
+        4980.6,
+        6006.3,
+        8683
+    ],
+    "KungFuMaster": [
+        7301,
+        15949,
+        34406,
+        14372,
+        36771
+    ],
+    "MsPacman": [
+        1251.6,
+        941.5,
+        590.3,
+        770.8,
+        1441.1
+    ],
+    "Pong": [
+        19.3,
+        18.5,
+        16.6,
+        6.2,
+        12.4
+    ],
+    "PrivateEye": [
+        100,
+        100,
+        100,
+        100,
+        100
+    ],
+    "Qbert": [
+        853,
+        575,
+        790.2,
+        1027.2,
+        483.2
+    ],
+    "RoadRunner": [
+        11355,
+        2033,
+        17106,
+        7614,
+        9965
+    ],
+    "Seaquest": [
+        752.2,
+        855.8,
+        691,
+        386,
+        621.6
+    ],
+    "UpNDown": [
+        3793,
+        5367.5,
+        2943.7,
+        3637.9,
+        1988.9
+    ]
+}
@@ -0,0 +1 @@
+{"Alien": [530.0], "Amidar": [38.8], "Assault": [500.1], "Asterix": [1734.0], "BankHeist": [192.5], "BattleZone": [7687.5], "Boxing": [15.1], "Breakout": [48.0], "ChopperCommand": [1350.0], "CrazyClimber": [56937.0], "DemonAttack": [3527.0], "Freeway": [21.8], "Frostbite": [255.0], "Gopher": [1256.0], "Hero": [3095.0], "Jamesbond": [87.5], "Kangaroo": [62.5], "Krull": [4890.8], "KungFuMaster": [18813.0], "MsPacman": [1265.6], "Pong": [-6.7], "PrivateEye": [56.3], "Qbert": [3952.0], "RoadRunner": [2500.0], "Seaquest": [208.0], "UpNDown": [2896.9]}
@@ -0,0 +1 @@
+{"Alien": [227.8], "Amidar": [5.8], "Assault": [222.4], "Asterix": [210.0], "BankHeist": [14.2], "BattleZone": [2360.0], "Boxing": [0.1], "Breakout": [1.7], "ChopperCommand": [811.0], "CrazyClimber": [10780.5], "DemonAttack": [152.1], "Freeway": [0.0], "Frostbite": [65.2], "Gopher": [257.6], "Hero": [1027.0], "Jamesbond": [29.0], "Kangaroo": [52.0], "Krull": [1598.0], "KungFuMaster": [258.5], "MsPacman": [307.3], "Pong": [-20.7], "PrivateEye": [24.9], "Qbert": [163.9], "RoadRunner": [11.5], "Seaquest": [68.4], "UpNDown": [533.4]}
@@ -0,0 +1 @@
+{"Alien": [677.1875, 435.46875, 515.9375, 425.3125, 1030.46875], "Amidar": [93.65625, 102.703125, 63.71875, 31.0, 80.203125], "Assault": [457.890625, 606.1875, 509.0625, 671.078125, 391.78125], "Asterix": [1061.71875, 1108.59375, 920.3125, 1065.625, 1485.15625], "BankHeist": [43.90625, 77.5, 0.0, 31.875, 17.96875], "BattleZone": [4031.2, 4031.2, 5359.375, 3484.375, 3250.0], "Boxing": [12.65625, 20.984375, 0.09375, -4.359375, 9.375], "Breakout": [10.3125, 16.4, 22.765625, 16.4, 16.046875], "ChopperCommand": [962.5, 896.875, 909.375, 848.4375, 1279.6875], "CrazyClimber": [54370.3125, 57220.3125, 51150.0, 87593.75, 62583.6], "DemonAttack": [288.90625, 208.1, 161.25, 177.578125, 204.84375], "Freeway": [0.0, 0.015625, 21.546875, 29.890625, 31.96875], "Frostbite": [262.34375, 248.75, 205.78125, 266.875, 200.625], "Gopher": [447.5, 732.5, 525.3125, 432.8125, 845.625], "Hero": [3061.5625, 3020.859375, 2060.390625, 2935.0, 2205.234375], "Jamesbond": [141.40625, 50.0, 110.9375, 123.4375, 76.5625], "Kangaroo": [65.625, 62.5, 25.0, 40.625, 62.5], "Krull": [3317.546875, 1933.71875, 1739.03125, 1379.625, 2654.28125], "KungFuMaster": [14318.75, 15023.4375, 14134.375, 21054.6875, 9781.25], "MsPacman": [1551.09375, 1221.25, 1199.0625, 1525.0, 1903.4375], "Pong": [21.0, -17.890625, 19.3125, 20.671875, 20.984375], "PrivateEye": [100.0, 0.0, 100.0, -25.140625, 0.0], "Qbert": [4258.984375, 150.0, 580.078125, 798.828125, 656.25], "RoadRunner": [3081.25, 900.0, 5906.25, 11176.5625, 7139.0625], "Seaquest": [711.5625, 455.3125, 565.625, 854.0625, 830.0], "UpNDown": [8614.53125, 1616.09375, 498.59375, 661.25, 5360.9375]}
@@ -0,0 +1,57 @@
+#! python
+
+import argparse
+from datetime import datetime
+from pathlib import Path
+import subprocess
+
+from omegaconf import OmegaConf
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-n', '--num-episodes', type=int, default=100, help='Number of evaluation episodes to collect.')
+    parser.add_argument('-p', '--num-envs', type=int, default=25, help='Number of environments used to collect the evaluation episodes.')
+    args = parser.parse_args()
+
+    path_to_config = Path('config') / 'trainer.yaml'
+    cfg = OmegaConf.load(path_to_config)
+
+    path_to_checkpoint = Path('checkpoints') / 'last.pt'
+    assert path_to_checkpoint.is_file()
+
+    cmd = f'python src/main.py hydra.run.dir=eval_outputs/{datetime.now().strftime("%Y-%m-%d/%H-%M-%S")} '
+
+    cmd += 'wandb.mode=online '
+    cmd += f'wandb.name=eval-{cfg.wandb.name} '
+    cmd += f'wandb.group=eval-{cfg.wandb.group} '
+
+    cmd += f'initialization.path_to_checkpoint={str(path_to_checkpoint.absolute())} '
+    cmd += 'initialization.load_tokenizer=True '
+    cmd += 'initialization.load_world_model=False '
+    cmd += 'initialization.load_actor_critic=True '
+
+    cmd += 'common.epochs=1 '
+    cmd += 'common.device=cuda:0 '
+    cmd += 'common.do_checkpoint=False '
+    cmd += 'common.seed=0 '
+
+    cmd += 'collection.test.num_episodes_to_save=0 '
+    cmd += f'collection.test.num_envs={args.num_envs} '
+    cmd += f'collection.test.config.num_episodes={args.num_episodes} '
+
+    cmd += 'training.should=False '
+
+    # Turn on data collection only
+    cmd += 'evaluation.should=True '
+    cmd += 'evaluation.every=1 '
+    cmd += 'evaluation.tokenizer.start_after_epochs=1 '
+    cmd += 'evaluation.tokenizer.save_reconstructions=False '
+    cmd += 'evaluation.world_model.start_after_epochs=1 '
+    cmd += 'evaluation.actor_critic.start_after_epochs=1 '
+
+    subprocess.run(cmd, shell=True, check=True)
+
+
+if __name__ == '__main__':
+    main()
@@ -0,0 +1,60 @@
+#! /usr/bin/python3
+
+import argparse
+from pathlib import Path
+import shutil
+import subprocess
+import yaml
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument('src_folder', type=str, help='Folder to import. Must be formatted as: USERNAME@HOSTNAME:PATH_TO_OUTPUTS_FOLDER/DATE/TIME')
+    parser.add_argument('-k', '--from-key', action='store_true')
+    parser.add_argument('-v', '--verbose', action='store_true')
+    args = parser.parse_args()
+
+    if args.from_key:
+        key_file = Path('path_runs.yaml')
+        assert key_file.is_file()
+        with key_file.open('r') as f:
+            runs = yaml.load(f, Loader=yaml.loader.SafeLoader)
+        src_folder = runs[args.src_folder]
+    else:
+        src_folder = args.src_folder
+
+    day, time = src_folder.split('/')[-2:]
+    host = src_folder.split('@')[1].split(':')[0]
+    dst_folder = Path(host) / day / time
+    print(dst_folder)
+    if dst_folder.is_dir():
+        if input(f'{dst_folder} exists, remove it ? [Y/n] ').lower() != 'n':
+            shutil.rmtree(dst_folder)
+        else:
+            print('Bye.')
+            return
+
+    dst_folder.mkdir(exist_ok=False, parents=True)
+
+    # Make symbolic link from key to folder
+    if args.from_key:
+        subprocess.run(f'ln -s {str(dst_folder)} {args.src_folder}', shell=True, check=True)
+
+    download_media = args.verbose and input('Download media/ ? [y/N] ').lower() == 'y'
+    download_last = not args.verbose or input('Download checkpoints/last.pt ? [Y/n] ').lower() != 'n'
+
+    folders = ['src', 'config', 'scripts']
+    if download_media:
+        folders.append('media')
+
+    for folder in folders:
+        subprocess.run(f'scp -r {src_folder}/{folder} {dst_folder}', shell=True, check=True)
+
+    if download_last:
+        checkpoint_folder = dst_folder / 'checkpoints'
+        checkpoint_folder.mkdir(exist_ok=False, parents=False)
+        subprocess.run(f'scp {src_folder}/checkpoints/last.pt {checkpoint_folder}', shell=True, check=True)
+
+
+if __name__ == '__main__':
+    main()
@@ -0,0 +1,29 @@
+#!/usr/bin/bash
+
+fps=15
+header=1
+mode="episode_replay"
+
+while [ "$1" != "" ]; do
+    case $1 in
+        -f | --fps )
+            shift
+            fps=$1
+            ;;
+        -h | --no-header )
+            header=0
+            ;;
+        -w | --world-model )
+            mode="world_model"
+            ;;
+        -a | --agent )
+            mode="agent"
+            ;;
+        * )
+            echo Invalid usage : $1
+            exit 1
+    esac
+    shift
+done
+
+python src/play.py hydra.run.dir=. hydra.output_subdir=null +mode="${mode}" +fps="${fps}" +header="${header}"
@@ -0,0 +1 @@
+python src/main.py common.resume=True hydra.output_subdir=null hydra.run.dir=.
@@ -0,0 +1,37 @@
+from pathlib import Path
+
+import torch
+from torch.distributions.categorical import Categorical
+import torch.nn as nn
+
+from models.actor_critic import ActorCritic
+from models.tokenizer import Tokenizer
+from models.world_model import WorldModel
+from utils import extract_state_dict
+
+
+class Agent(nn.Module):
+    def __init__(self, tokenizer: Tokenizer, world_model: WorldModel, actor_critic: ActorCritic):
+        super().__init__()
+        self.tokenizer = tokenizer
+        self.world_model = world_model
+        self.actor_critic = actor_critic
+
+    @property
+    def device(self):
+        return self.actor_critic.conv1.weight.device
+
+    def load(self, path_to_checkpoint: Path, device: torch.device, load_tokenizer: bool = True, load_world_model: bool = True, load_actor_critic: bool = True) -> None:
+        agent_state_dict = torch.load(path_to_checkpoint, map_location=device)
+        if load_tokenizer:
+            self.tokenizer.load_state_dict(extract_state_dict(agent_state_dict, 'tokenizer'))
+        if load_world_model:
+            self.world_model.load_state_dict(extract_state_dict(agent_state_dict, 'world_model'))
+        if load_actor_critic:
+            self.actor_critic.load_state_dict(extract_state_dict(agent_state_dict, 'actor_critic'))
+
+    def act(self, obs: torch.FloatTensor, should_sample: bool = True, temperature: float = 1.0) -> torch.LongTensor:
+        input_ac = obs if self.actor_critic.use_original_obs else torch.clamp(self.tokenizer.encode_decode(obs, should_preprocess=True, should_postprocess=True), 0, 1)
+        logits_actions = self.actor_critic(input_ac).logits_actions[:, -1] / temperature
+        act_token = Categorical(logits=logits_actions).sample() if should_sample else logits_actions.argmax(dim=-1)
+        return act_token
@@ -0,0 +1,125 @@
+import random
+import sys
+from typing import List, Optional, Union
+
+from einops import rearrange
+import numpy as np
+import torch
+from tqdm import tqdm
+import wandb
+
+from agent import Agent
+from dataset import EpisodesDataset
+from envs import SingleProcessEnv, MultiProcessEnv
+from episode import Episode
+from utils import EpisodeDirManager, RandomHeuristic
+
+
+class Collector:
+    def __init__(self, env: Union[SingleProcessEnv, MultiProcessEnv], dataset: EpisodesDataset, episode_dir_manager: EpisodeDirManager) -> None:
+        self.env = env
+        self.dataset = dataset
+        self.episode_dir_manager = episode_dir_manager
+        self.obs = self.env.reset()
+        self.episode_ids = [None] * self.env.num_envs
+        self.heuristic = RandomHeuristic(self.env.num_actions)
+
+    @torch.no_grad()
+    def collect(self, agent: Agent, epoch: int, epsilon: float, should_sample: bool, temperature: float, burn_in: int, *, num_steps: Optional[int] = None, num_episodes: Optional[int] = None):
+        assert self.env.num_actions == agent.world_model.act_vocab_size
+        assert 0 <= epsilon <= 1
+
+        assert (num_steps is None) != (num_episodes is None)
+        should_stop = lambda steps, episodes: steps >= num_steps if num_steps is not None else episodes >= num_episodes
+
+        to_log = []
+        steps, episodes = 0, 0
+        returns = []
+        observations, actions, rewards, dones = [], [], [], []
+
+        burnin_obs_rec, mask_padding = None, None
+        if set(self.episode_ids) != {None} and burn_in > 0:
+            current_episodes = [self.dataset.get_episode(episode_id) for episode_id in self.episode_ids]
+            segmented_episodes = [episode.segment(start=len(episode) - burn_in, stop=len(episode), should_pad=True) for episode in current_episodes]
+            mask_padding = torch.stack([episode.mask_padding for episode in segmented_episodes], dim=0).to(agent.device)
+            burnin_obs = torch.stack([episode.observations for episode in segmented_episodes], dim=0).float().div(255).to(agent.device)
+            burnin_obs_rec = torch.clamp(agent.tokenizer.encode_decode(burnin_obs, should_preprocess=True, should_postprocess=True), 0, 1)
+
+        agent.actor_critic.reset(n=self.env.num_envs, burnin_observations=burnin_obs_rec, mask_padding=mask_padding)
+        pbar = tqdm(total=num_steps if num_steps is not None else num_episodes, desc=f'Experience collection ({self.dataset.name})', file=sys.stdout)
+
+        while not should_stop(steps, episodes):
+
+            observations.append(self.obs)
+            obs = rearrange(torch.FloatTensor(self.obs).div(255), 'n h w c -> n c h w').to(agent.device)
+            act = agent.act(obs, should_sample=should_sample, temperature=temperature).cpu().numpy()
+
+            if random.random() < epsilon:
+                act = self.heuristic.act(obs).cpu().numpy()
+
+            self.obs, reward, done, _ = self.env.step(act)
+
+            actions.append(act)
+            rewards.append(reward)
+            dones.append(done)
+
+            new_steps = len(self.env.mask_new_dones)
+            steps += new_steps
+            pbar.update(new_steps if num_steps is not None else 0)
+
+            # Warning: with EpisodicLifeEnv + MultiProcessEnv, reset is ignored if not a real done.
+            # Thus, segments of experience following a life loss and preceding a general done are discarded.
+            # Not a problem with a SingleProcessEnv.
+
+            if self.env.should_reset():
+                self.add_experience_to_dataset(observations, actions, rewards, dones)
+
+                new_episodes = self.env.num_envs
+                episodes += new_episodes
+                pbar.update(new_episodes if num_episodes is not None else 0)
+
+                for episode_id in self.episode_ids:
+                    episode = self.dataset.get_episode(episode_id)
+                    self.episode_dir_manager.save(episode, episode_id, epoch)
+                    metrics_episode = {k: v for k, v in episode.compute_metrics().__dict__.items()}
+                    metrics_episode['episode_num'] = episode_id
+                    metrics_episode['action_histogram'] = wandb.Histogram(np_histogram=np.histogram(episode.actions.numpy(), bins=np.arange(0, self.env.num_actions + 1) - 0.5, density=True))
+                    to_log.append({f'{self.dataset.name}/{k}': v for k, v in metrics_episode.items()})
+                    returns.append(metrics_episode['episode_return'])
+
+                self.obs = self.env.reset()
+                self.episode_ids = [None] * self.env.num_envs
+                agent.actor_critic.reset(n=self.env.num_envs)
+                observations, actions, rewards, dones = [], [], [], []
+
+        # Add incomplete episodes to dataset, and complete them later.
+        if len(observations) > 0:
+            self.add_experience_to_dataset(observations, actions, rewards, dones)
+
+        agent.actor_critic.clear()
+
+        metrics_collect = {
+            '#episodes': len(self.dataset),
+            '#steps': sum(map(len, self.dataset.episodes)),
+        }
+        if len(returns) > 0:
+            metrics_collect['return'] = np.mean(returns)
+        metrics_collect = {f'{self.dataset.name}/{k}': v for k, v in metrics_collect.items()}
+        to_log.append(metrics_collect)
+
+        return to_log
+
+    def add_experience_to_dataset(self, observations: List[np.ndarray], actions: List[np.ndarray], rewards: List[np.ndarray], dones: List[np.ndarray]) -> None:
+        assert len(observations) == len(actions) == len(rewards) == len(dones)
+        for i, (o, a, r, d) in enumerate(zip(*map(lambda arr: np.swapaxes(arr, 0, 1), [observations, actions, rewards, dones]))):  # Make everything (N, T, ...) instead of (T, N, ...)
+            episode = Episode(
+                observations=torch.ByteTensor(o).permute(0, 3, 1, 2).contiguous(),  # channel-first
+                actions=torch.LongTensor(a),
+                rewards=torch.FloatTensor(r),
+                ends=torch.LongTensor(d),
+                mask_padding=torch.ones(d.shape[0], dtype=torch.bool),
+            )
+            if self.episode_ids[i] is None:
+                self.episode_ids[i] = self.dataset.add_episode(episode)
+            else:
+                self.dataset.update_episode(self.episode_ids[i], episode)
@@ -0,0 +1,163 @@
+from collections import deque
+import math
+from pathlib import Path
+import random
+from typing import Dict, List, Optional, Tuple
+
+import psutil
+import torch
+
+from episode import Episode
+
+Batch = Dict[str, torch.Tensor]
+
+
+class EpisodesDataset:
+    def __init__(self, max_num_episodes: Optional[int] = None, name: Optional[str] = None) -> None:
+        self.max_num_episodes = max_num_episodes
+        self.name = name if name is not None else 'dataset'
+        self.num_seen_episodes = 0
+        self.episodes = deque()
+        self.episode_id_to_queue_idx = dict()
+        self.newly_modified_episodes, self.newly_deleted_episodes = set(), set()
+
+    def __len__(self) -> int:
+        return len(self.episodes)
+
+    def clear(self) -> None:
+        self.episodes = deque()
+        self.episode_id_to_queue_idx = dict()
+
+    def add_episode(self, episode: Episode) -> int:
+        if self.max_num_episodes is not None and len(self.episodes) == self.max_num_episodes:
+            self._popleft()
+        episode_id = self._append_new_episode(episode)
+        return episode_id
+
+    def get_episode(self, episode_id: int) -> Episode:
+        assert episode_id in self.episode_id_to_queue_idx
+        queue_idx = self.episode_id_to_queue_idx[episode_id]
+        return self.episodes[queue_idx]
+
+    def update_episode(self, episode_id: int, new_episode: Episode) -> None:
+        assert episode_id in self.episode_id_to_queue_idx
+        queue_idx = self.episode_id_to_queue_idx[episode_id]
+        merged_episode = self.episodes[queue_idx].merge(new_episode)
+        self.episodes[queue_idx] = merged_episode
+        self.newly_modified_episodes.add(episode_id)
+
+    def _popleft(self) -> Episode:
+        id_to_delete = [k for k, v in self.episode_id_to_queue_idx.items() if v == 0]
+        assert len(id_to_delete) == 1
+        self.newly_deleted_episodes.add(id_to_delete[0])
+        self.episode_id_to_queue_idx = {k: v - 1 for k, v in self.episode_id_to_queue_idx.items() if v > 0}
+        return self.episodes.popleft()
+
+    def _append_new_episode(self, episode):
+        episode_id = self.num_seen_episodes
+        self.episode_id_to_queue_idx[episode_id] = len(self.episodes)
+        self.episodes.append(episode)
+        self.num_seen_episodes += 1
+        self.newly_modified_episodes.add(episode_id)
+        return episode_id
+
+    def sample_batch(self, batch_num_samples: int, sequence_length: int, weights: Optional[Tuple[float]] = None, sample_from_start: bool = True) -> Batch:
+        return self._collate_episodes_segments(self._sample_episodes_segments(batch_num_samples, sequence_length, weights, sample_from_start))
+
+    def _sample_episodes_segments(self, batch_num_samples: int, sequence_length: int, weights: Optional[Tuple[float]], sample_from_start: bool) -> List[Episode]:
+        num_episodes = len(self.episodes)
+        num_weights = len(weights) if weights is not None else 0
+
+        if num_weights < num_episodes:
+            weights = [1] * num_episodes
+        else:
+            assert all([0 <= x <= 1 for x in weights]) and sum(weights) == 1
+            sizes = [num_episodes // num_weights + (num_episodes % num_weights) * (i == num_weights - 1) for i in range(num_weights)]
+            weights = [w / s for (w, s) in zip(weights, sizes) for _ in range(s)]
+
+        sampled_episodes = random.choices(self.episodes, k=batch_num_samples, weights=weights)
+
+        sampled_episodes_segments = []
+        for sampled_episode in sampled_episodes:
+            if sample_from_start:
+                start = random.randint(0, len(sampled_episode) - 1)
+                stop = start + sequence_length
+            else:
+                stop = random.randint(1, len(sampled_episode))
+                start = stop - sequence_length
+            sampled_episodes_segments.append(sampled_episode.segment(start, stop, should_pad=True))
+            assert len(sampled_episodes_segments[-1]) == sequence_length
+        return sampled_episodes_segments
+
+    def _collate_episodes_segments(self, episodes_segments: List[Episode]) -> Batch:
+        episodes_segments = [e_s.__dict__ for e_s in episodes_segments]
+        batch = {}
+        for k in episodes_segments[0]:
+            batch[k] = torch.stack([e_s[k] for e_s in episodes_segments])
+        batch['observations'] = batch['observations'].float() / 255.0  # int8 to float and scale
+        return batch
+
+    def traverse(self, batch_num_samples: int, chunk_size: int):
+        for episode in self.episodes:
+            chunks = [episode.segment(start=i * chunk_size, stop=(i + 1) * chunk_size, should_pad=True) for i in range(math.ceil(len(episode) / chunk_size))]
+            batches = [chunks[i * batch_num_samples: (i + 1) * batch_num_samples] for i in range(math.ceil(len(chunks) / batch_num_samples))]
+            for b in batches:
+                yield self._collate_episodes_segments(b)
+
+    def update_disk_checkpoint(self, directory: Path) -> None:
+        assert directory.is_dir()
+        for episode_id in self.newly_modified_episodes:
+            episode = self.get_episode(episode_id)
+            episode.save(directory / f'{episode_id}.pt')
+        for episode_id in self.newly_deleted_episodes:
+            (directory / f'{episode_id}.pt').unlink()
+        self.newly_modified_episodes, self.newly_deleted_episodes = set(), set()
+
+    def load_disk_checkpoint(self, directory: Path) -> None:
+        assert directory.is_dir() and len(self.episodes) == 0
+        episode_ids = sorted([int(p.stem) for p in directory.iterdir()])
+        self.num_seen_episodes = episode_ids[-1] + 1
+        for episode_id in episode_ids:
+            episode = Episode(**torch.load(directory / f'{episode_id}.pt'))
+            self.episode_id_to_queue_idx[episode_id] = len(self.episodes)
+            self.episodes.append(episode)
+
+
+class EpisodesDatasetRamMonitoring(EpisodesDataset):
+    """
+    Prevent episode dataset from going out of RAM.
+    Warning: % looks at system wide RAM usage while G looks only at process RAM usage.
+    """
+    def __init__(self, max_ram_usage: str, name: Optional[str] = None) -> None:
+        super().__init__(max_num_episodes=None, name=name)
+        self.max_ram_usage = max_ram_usage
+        self.num_steps = 0
+        self.max_num_steps = None
+
+        max_ram_usage = str(max_ram_usage)
+        if max_ram_usage.endswith('%'):
+            m = int(max_ram_usage.split('%')[0])
+            assert 0 < m < 100
+            self.check_ram_usage = lambda: psutil.virtual_memory().percent > m
+        else:
+            assert max_ram_usage.endswith('G')
+            m = float(max_ram_usage.split('G')[0])
+            self.check_ram_usage = lambda: psutil.Process().memory_info()[0] / 2 ** 30 > m
+
+    def clear(self) -> None:
+        super().clear()
+        self.num_steps = 0
+
+    def add_episode(self, episode: Episode) -> int:
+        if self.max_num_steps is None and self.check_ram_usage():
+            self.max_num_steps = self.num_steps
+        self.num_steps += len(episode)
+        while (self.max_num_steps is not None) and (self.num_steps > self.max_num_steps):
+            self._popleft()
+        episode_id = self._append_new_episode(episode)
+        return episode_id
+
+    def _popleft(self) -> Episode:
+        episode = super()._popleft()
+        self.num_steps -= len(episode)
+        return episode
@@ -0,0 +1,4 @@
+from .multi_process_env import MultiProcessEnv
+from .wrappers import make_atari, ResizeObsWrapper
+from .single_process_env import SingleProcessEnv
+from .world_model_env import WorldModelEnv
@@ -0,0 +1,27 @@
+import numpy as np
+
+
+class DoneTrackerEnv:
+    def __init__(self, num_envs: int) -> None:
+        """Monitor env dones: 0 when not done, 1 when done, 2 when already done."""
+        self.num_envs = num_envs
+        self.done_tracker = None
+        self.reset_done_tracker()
+
+    def reset_done_tracker(self) -> None:
+        self.done_tracker = np.zeros(self.num_envs, dtype=np.uint8)
+
+    def update_done_tracker(self, done: np.ndarray) -> None:
+        self.done_tracker = np.clip(2 * self.done_tracker + done, 0, 2)
+
+    @property
+    def num_envs_done(self) -> int:
+        return (self.done_tracker > 0).sum()
+
+    @property
+    def mask_dones(self) -> np.ndarray:
+        return np.logical_not(self.done_tracker)
+
+    @property
+    def mask_new_dones(self) -> np.ndarray:
+        return np.logical_not(self.done_tracker[self.done_tracker <= 1])
@@ -0,0 +1,94 @@
+from dataclasses import astuple, dataclass
+from enum import Enum
+from multiprocessing import Pipe, Process
+from multiprocessing.connection import Connection
+from typing import Any, Callable, Iterator, List, Optional, Tuple
+
+import numpy as np
+
+from .done_tracker import DoneTrackerEnv
+
+
+class MessageType(Enum):
+    RESET = 0
+    RESET_RETURN = 1
+    STEP = 2
+    STEP_RETURN = 3
+    CLOSE = 4
+
+
+@dataclass
+class Message:
+    type: MessageType
+    content: Optional[Any] = None
+
+    def __iter__(self) -> Iterator:
+        return iter(astuple(self))
+
+
+def child_env(child_id: int, env_fn: Callable, child_conn: Connection) -> None:
+    np.random.seed(child_id + np.random.randint(0, 2 ** 31 - 1))
+    env = env_fn()
+    while True:
+        message_type, content = child_conn.recv()
+        if message_type == MessageType.RESET:
+            obs = env.reset()
+            child_conn.send(Message(MessageType.RESET_RETURN, obs))
+        elif message_type == MessageType.STEP:
+            obs, rew, done, _ = env.step(content)
+            if done:
+                obs = env.reset()
+            child_conn.send(Message(MessageType.STEP_RETURN, (obs, rew, done, None)))
+        elif message_type == MessageType.CLOSE:
+            child_conn.close()
+            return
+        else:
+            raise NotImplementedError
+
+
+class MultiProcessEnv(DoneTrackerEnv):
+    def __init__(self, env_fn: Callable, num_envs: int, should_wait_num_envs_ratio: float) -> None:
+        super().__init__(num_envs)
+        self.num_actions = env_fn().env.action_space.n
+        self.should_wait_num_envs_ratio = should_wait_num_envs_ratio
+        self.processes, self.parent_conns = [], []
+        for child_id in range(num_envs):
+            parent_conn, child_conn = Pipe()
+            self.parent_conns.append(parent_conn)
+            p = Process(target=child_env, args=(child_id, env_fn, child_conn), daemon=True)
+            self.processes.append(p)
+        for p in self.processes:
+            p.start()
+
+    def should_reset(self) -> bool:
+        return (self.num_envs_done / self.num_envs) >= self.should_wait_num_envs_ratio
+
+    def _receive(self, check_type: Optional[MessageType] = None) -> List[Any]:
+        messages = [parent_conn.recv() for parent_conn in self.parent_conns]
+        if check_type is not None:
+            assert all([m.type == check_type for m in messages])
+        return [m.content for m in messages]
+
+    def reset(self) -> np.ndarray:
+        self.reset_done_tracker()
+        for parent_conn in self.parent_conns:
+            parent_conn.send(Message(MessageType.RESET))
+        content = self._receive(check_type=MessageType.RESET_RETURN)
+        return np.stack(content)
+
+    def step(self, actions: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Any]:
+        for parent_conn, action in zip(self.parent_conns, actions):
+            parent_conn.send(Message(MessageType.STEP, action))
+        content = self._receive(check_type=MessageType.STEP_RETURN)
+        obs, rew, done, _ = zip(*content)
+        done = np.stack(done)
+        self.update_done_tracker(done)
+        return np.stack(obs), np.stack(rew), done, None
+
+    def close(self) -> None:
+        for parent_conn in self.parent_conns:
+            parent_conn.send(Message(MessageType.CLOSE))
+        for p in self.processes:
+            p.join()
+        for parent_conn in self.parent_conns:
+            parent_conn.close()
@@ -0,0 +1,32 @@
+from typing import Any, Tuple
+
+import numpy as np
+
+from .done_tracker import DoneTrackerEnv
+
+
+class SingleProcessEnv(DoneTrackerEnv):
+    def __init__(self, env_fn):
+        super().__init__(num_envs=1)
+        self.env = env_fn()
+        self.num_actions = self.env.action_space.n
+
+    def should_reset(self) -> bool:
+        return self.num_envs_done == 1
+
+    def reset(self) -> np.ndarray:
+        self.reset_done_tracker()
+        obs = self.env.reset()
+        return obs[None, ...]
+
+    def step(self, action) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Any]:
+        obs, reward, done, _ = self.env.step(action[0])  # action is supposed to be ndarray (1,)
+        done = np.array([done])
+        self.update_done_tracker(done)
+        return obs[None, ...], np.array([reward]), done, None
+
+    def render(self) -> None:
+        self.env.render()
+
+    def close(self) -> None:
+        self.env.close()
@@ -0,0 +1,113 @@
+import random
+from typing import List, Optional, Union
+
+import gym
+from einops import rearrange
+import numpy as np
+from PIL import Image
+import torch
+from torch.distributions.categorical import Categorical
+import torchvision
+
+from utils import extract_state_dict
+
+
+class WorldModelEnv:
+
+    def __init__(self, tokenizer: torch.nn.Module, world_model: torch.nn.Module, device: Union[str, torch.device], pretrained_agent_path: Optional[str] = None, env: Optional[gym.Env] = None) -> None:
+
+        self.device = torch.device(device)
+        self.world_model = world_model.to(self.device).eval()
+        self.tokenizer = tokenizer.to(self.device).eval()
+
+        if pretrained_agent_path is not None:
+            agent_state_dict = torch.load(pretrained_agent_path)
+            self.world_model.load_state_dict(extract_state_dict(agent_state_dict, 'world_model'))
+            incompatible_keys = self.tokenizer.load_state_dict(extract_state_dict(agent_state_dict, 'tokenizer'), strict=False)
+            assert not incompatible_keys.missing_keys
+            assert (not incompatible_keys.unexpected_keys) or all([k.startswith('lpips.') for k in incompatible_keys.unexpected_keys])
+
+        self.keys_values_wm, self.obs_tokens, self._num_observations_tokens = None, None, None
+
+        self.env = env
+
+    @property
+    def num_observations_tokens(self) -> int:
+        return self._num_observations_tokens
+
+    @torch.no_grad()
+    def reset(self) -> torch.FloatTensor:
+        assert self.env is not None
+        obs = torchvision.transforms.functional.to_tensor(self.env.reset()).to(self.device).unsqueeze(0)  # (1, C, H, W) in [0., 1.]
+        return self.reset_from_initial_observations(obs)
+
+    @torch.no_grad()
+    def reset_from_initial_observations(self, observations: torch.FloatTensor) -> torch.FloatTensor:
+        obs_tokens = self.tokenizer.encode(observations, should_preprocess=True).tokens    # (B, C, H, W) -> (B, K)
+        _, num_observations_tokens = obs_tokens.shape
+        if self.num_observations_tokens is None:
+            self._num_observations_tokens = num_observations_tokens
+
+        _ = self.refresh_keys_values_with_initial_obs_tokens(obs_tokens)
+        self.obs_tokens = obs_tokens
+
+        return self.decode_obs_tokens()
+
+    @torch.no_grad()
+    def refresh_keys_values_with_initial_obs_tokens(self, obs_tokens: torch.LongTensor) -> torch.FloatTensor:
+        n, num_observations_tokens = obs_tokens.shape
+        assert num_observations_tokens == self.num_observations_tokens
+        self.keys_values_wm = self.world_model.transformer.generate_empty_keys_values(n=n, max_tokens=self.world_model.config.max_tokens)
+        outputs_wm = self.world_model(obs_tokens, past_keys_values=self.keys_values_wm)
+        return outputs_wm.output_sequence  # (B, K, E)
+
+    @torch.no_grad()
+    def step(self, action: Union[int, np.ndarray, torch.LongTensor], should_predict_next_obs: bool = True) -> None:
+        assert self.keys_values_wm is not None and self.num_observations_tokens is not None
+
+        num_passes = 1 + self.num_observations_tokens if should_predict_next_obs else 1
+
+        output_sequence, obs_tokens = [], []
+
+        if self.keys_values_wm.size + num_passes > self.world_model.config.max_tokens:
+            _ = self.refresh_keys_values_with_initial_obs_tokens(self.obs_tokens)
+
+        token = action.clone().detach() if isinstance(action, torch.Tensor) else torch.tensor(action, dtype=torch.long)
+        token = token.reshape(-1, 1).to(self.device)  # (B, 1)
+
+        for k in range(num_passes):  # assumption that there is only one action token.
+
+            outputs_wm = self.world_model(token, past_keys_values=self.keys_values_wm)
+            output_sequence.append(outputs_wm.output_sequence)
+
+            if k == 0:
+                reward = Categorical(logits=outputs_wm.logits_rewards).sample().float().cpu().numpy().reshape(-1) - 1   # (B,)
+                done = Categorical(logits=outputs_wm.logits_ends).sample().cpu().numpy().astype(bool).reshape(-1)       # (B,)
+
+            if k < self.num_observations_tokens:
+                token = Categorical(logits=outputs_wm.logits_observations).sample()
+                obs_tokens.append(token)
+
+        output_sequence = torch.cat(output_sequence, dim=1)   # (B, 1 + K, E)
+        self.obs_tokens = torch.cat(obs_tokens, dim=1)        # (B, K)
+
+        obs = self.decode_obs_tokens() if should_predict_next_obs else None
+        return obs, reward, done, None
+
+    @torch.no_grad()
+    def render_batch(self) -> List[Image.Image]:
+        frames = self.decode_obs_tokens().detach().cpu()
+        frames = rearrange(frames, 'b c h w -> b h w c').mul(255).numpy().astype(np.uint8)
+        return [Image.fromarray(frame) for frame in frames]
+
+    @torch.no_grad()
+    def decode_obs_tokens(self) -> List[Image.Image]:
+        embedded_tokens = self.tokenizer.embedding(self.obs_tokens)     # (B, K, E)
+        z = rearrange(embedded_tokens, 'b (h w) e -> b e h w', h=int(np.sqrt(self.num_observations_tokens)))
+        rec = self.tokenizer.decode(z, should_postprocess=True)         # (B, C, H, W)
+        return torch.clamp(rec, 0, 1)
+
+    @torch.no_grad()
+    def render(self):
+        assert self.obs_tokens.shape == (1, self.num_observations_tokens)
+        return self.render_batch()[0]
@@ -0,0 +1,147 @@
+"""
+Credits to https://github.com/openai/baselines/blob/master/baselines/common/atari_wrappers.py
+"""
+
+from typing import Tuple
+
+import gym
+import numpy as np
+from PIL import Image
+
+
+def make_atari(id, size=64, max_episode_steps=None, noop_max=30, frame_skip=4, done_on_life_loss=False, clip_reward=False):
+    env = gym.make(id)
+    assert 'NoFrameskip' in env.spec.id or 'Frameskip' not in env.spec
+    env = ResizeObsWrapper(env, (size, size))
+    if clip_reward:
+        env = RewardClippingWrapper(env)
+    if max_episode_steps is not None:
+        env = gym.wrappers.TimeLimit(env, max_episode_steps=max_episode_steps)
+    if noop_max is not None:
+        env = NoopResetEnv(env, noop_max=noop_max)
+    env = MaxAndSkipEnv(env, skip=frame_skip)
+    if done_on_life_loss:
+        env = EpisodicLifeEnv(env)
+    return env
+
+
+class ResizeObsWrapper(gym.ObservationWrapper):
+    def __init__(self, env: gym.Env, size: Tuple[int, int]) -> None:
+        gym.ObservationWrapper.__init__(self, env)
+        self.size = tuple(size)
+        self.observation_space = gym.spaces.Box(low=0, high=255, shape=(size[0], size[1], 3), dtype=np.uint8)
+        self.original_obs = None
+
+    def resize(self, obs: np.ndarray):
+        img = Image.fromarray(obs)
+        img = img.resize(self.size, Image.BILINEAR)
+        return np.array(img)
+
+    def observation(self, observation: np.ndarray) -> np.ndarray:
+        self.original_obs = observation
+        return self.resize(observation)
+
+
+class RewardClippingWrapper(gym.RewardWrapper):
+    def reward(self, reward):
+        return np.sign(reward)
+
+
+class NoopResetEnv(gym.Wrapper):
+    def __init__(self, env, noop_max=30):
+        """Sample initial states by taking random number of no-ops on reset.
+        No-op is assumed to be action 0.
+        """
+        gym.Wrapper.__init__(self, env)
+        self.noop_max = noop_max
+        self.override_num_noops = None
+        self.noop_action = 0
+        assert env.unwrapped.get_action_meanings()[0] == 'NOOP'
+
+    def reset(self, **kwargs):
+        """ Do no-op action for a number of steps in [1, noop_max]."""
+        self.env.reset(**kwargs)
+        if self.override_num_noops is not None:
+            noops = self.override_num_noops
+        else:
+            noops = self.unwrapped.np_random.randint(1, self.noop_max + 1)
+        assert noops > 0
+        obs = None
+        for _ in range(noops):
+            obs, _, done, _ = self.env.step(self.noop_action)
+            if done:
+                obs = self.env.reset(**kwargs)
+        return obs
+
+    def step(self, action):
+        return self.env.step(action)
+
+
+class EpisodicLifeEnv(gym.Wrapper):
+    def __init__(self, env):
+        """Make end-of-life == end-of-episode, but only reset on true game over.
+        Done by DeepMind for the DQN and co. since it helps value estimation.
+        """
+        gym.Wrapper.__init__(self, env)
+        self.lives = 0
+        self.was_real_done = True
+
+    def step(self, action):
+        obs, reward, done, info = self.env.step(action)
+        self.was_real_done = done
+        # check current lives, make loss of life terminal,
+        # then update lives to handle bonus lives
+        lives = self.env.unwrapped.ale.lives()
+        if lives < self.lives and lives > 0:
+            # for Qbert sometimes we stay in lives == 0 condition for a few frames
+            # so it's important to keep lives > 0, so that we only reset once
+            # the environment advertises done.
+            done = True
+        self.lives = lives
+        return obs, reward, done, info
+
+    def reset(self, **kwargs):
+        """Reset only when lives are exhausted.
+        This way all states are still reachable even though lives are episodic,
+        and the learner need not know about any of this behind-the-scenes.
+        """
+        if self.was_real_done:
+            obs = self.env.reset(**kwargs)
+        else:
+            # no-op step to advance from terminal/lost life state
+            obs, _, _, _ = self.env.step(0)
+        self.lives = self.env.unwrapped.ale.lives()
+        return obs
+
+
+class MaxAndSkipEnv(gym.Wrapper):
+    def __init__(self, env, skip=4):
+        """Return only every `skip`-th frame"""
+        gym.Wrapper.__init__(self, env)
+        assert skip > 0
+        # most recent raw observations (for max pooling across time steps)
+        self._obs_buffer = np.zeros((2,) + env.observation_space.shape, dtype=np.uint8)
+        self._skip = skip
+        self.max_frame = np.zeros(env.observation_space.shape, dtype=np.uint8)
+
+    def step(self, action):
+        """Repeat action, sum reward, and max over last observations."""
+        total_reward = 0.0
+        done = None
+        for i in range(self._skip):
+            obs, reward, done, info = self.env.step(action)
+            if i == self._skip - 2:
+                self._obs_buffer[0] = obs
+            if i == self._skip - 1:
+                self._obs_buffer[1] = obs
+            total_reward += reward
+            if done:
+                break
+        # Note that the observation on the done=True frame
+        # doesn't matter
+        self.max_frame = self._obs_buffer.max(axis=0)
+
+        return self.max_frame, total_reward, done, info
+
+    def reset(self, **kwargs):
+        return self.env.reset(**kwargs)
@@ -0,0 +1,76 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+
+
+@dataclass
+class EpisodeMetrics:
+    episode_length: int
+    episode_return: float
+
+
+@dataclass
+class Episode:
+    observations: torch.ByteTensor
+    actions: torch.LongTensor
+    rewards: torch.FloatTensor
+    ends: torch.LongTensor
+    mask_padding: torch.BoolTensor
+
+    def __post_init__(self):
+        assert len(self.observations) == len(self.actions) == len(self.rewards) == len(self.ends) == len(self.mask_padding)
+        if self.ends.sum() > 0:
+            idx_end = torch.argmax(self.ends) + 1
+            self.observations = self.observations[:idx_end]
+            self.actions = self.actions[:idx_end]
+            self.rewards = self.rewards[:idx_end]
+            self.ends = self.ends[:idx_end]
+            self.mask_padding = self.mask_padding[:idx_end]
+
+    def __len__(self) -> int:
+        return self.observations.size(0)
+
+    def merge(self, other: Episode) -> Episode:
+        return Episode(
+            torch.cat((self.observations, other.observations), dim=0),
+            torch.cat((self.actions, other.actions), dim=0),
+            torch.cat((self.rewards, other.rewards), dim=0),
+            torch.cat((self.ends, other.ends), dim=0),
+            torch.cat((self.mask_padding, other.mask_padding), dim=0),
+        )
+
+    def segment(self, start: int, stop: int, should_pad: bool = False) -> Episode:
+        assert start < len(self) and stop > 0 and start < stop
+        padding_length_right = max(0, stop - len(self))
+        padding_length_left = max(0, -start)
+        assert padding_length_right == padding_length_left == 0 or should_pad
+
+        def pad(x):
+            pad_right = torch.nn.functional.pad(x, [0 for _ in range(2 * x.ndim - 1)] + [padding_length_right]) if padding_length_right > 0 else x
+            return torch.nn.functional.pad(pad_right, [0 for _ in range(2 * x.ndim - 2)] + [padding_length_left, 0]) if padding_length_left > 0 else pad_right
+
+        start = max(0, start)
+        stop = min(len(self), stop)
+        segment = Episode(
+            self.observations[start:stop],
+            self.actions[start:stop],
+            self.rewards[start:stop],
+            self.ends[start:stop],
+            self.mask_padding[start:stop],
+        )
+
+        segment.observations = pad(segment.observations)
+        segment.actions = pad(segment.actions)
+        segment.rewards = pad(segment.rewards)
+        segment.ends = pad(segment.ends)
+        segment.mask_padding = torch.cat((torch.zeros(padding_length_left, dtype=torch.bool), segment.mask_padding, torch.zeros(padding_length_right, dtype=torch.bool)), dim=0)
+
+        return segment
+
+    def compute_metrics(self) -> EpisodeMetrics:
+        return EpisodeMetrics(len(self), self.rewards.sum())
+
+    def save(self, path: Path) -> None:
+        torch.save(self.__dict__, path)
@@ -0,0 +1,3 @@
+from .agent_env import AgentEnv
+from .episode_replay_env import EpisodeReplayEnv
+from .game import Game
@@ -0,0 +1,54 @@
+from einops import rearrange
+import numpy as np
+from PIL import Image
+import torch
+
+from agent import Agent
+from envs import SingleProcessEnv
+from game.keymap import get_keymap_and_action_names
+
+
+class AgentEnv:
+    def __init__(self, agent: Agent, env: SingleProcessEnv, keymap_name: str) -> None:
+        assert isinstance(env, SingleProcessEnv)
+        self.agent = agent
+        self.env = env
+        _, self.action_names = get_keymap_and_action_names(keymap_name)
+        self.obs = None
+        self._t = None
+        self._return = None
+
+    def _to_tensor(self, obs: np.ndarray):
+        return rearrange(torch.FloatTensor(obs).div(255), 'n h w c -> n c h w').to(self.agent.device)
+
+    def _to_array(self, obs: torch.FloatTensor):
+        assert obs.ndim == 4 and obs.size(0) == 1
+        return obs[0].mul(255).permute(1, 2, 0).cpu().numpy().astype(np.uint8)
+
+    def reset(self):
+        obs = self.env.reset()
+        self.obs = self._to_tensor(obs)
+        self.agent.actor_critic.reset(1)
+        self._t = 0
+        self._return = 0
+        return obs
+
+    def step(self, *args, **kwargs) -> torch.FloatTensor:
+        with torch.no_grad():
+            act = self.agent.act(self.obs, should_sample=True).cpu().numpy()
+        obs, reward, done, _ = self.env.step(act)
+        self.obs = self._to_tensor(obs)
+        self._t += 1
+        self._return += reward[0]
+        info = {
+            'timestep': self._t,
+            'action': self.action_names[act[0]],
+            'return': self._return,
+        }
+        return obs, reward, done, info
+
+    def render(self) -> Image.Image:
+        assert self.obs.size() == (1, 3, 64, 64)
+        rec = torch.clamp(self.agent.tokenizer.encode_decode(self.obs, should_preprocess=True, should_postprocess=True), 0, 1)
+        arr = self._to_array(torch.cat((self.obs, rec), dim=-1))
+        return Image.fromarray(arr)
@@ -0,0 +1,111 @@
+from pathlib import Path
+
+import numpy as np
+from PIL import Image
+import torch
+
+from episode import Episode
+from game.keymap import get_keymap_and_action_names
+
+
+class EpisodeReplayEnv:
+    def __init__(self, replay_keymap_name: str, episode_dir: Path) -> None:
+        _, self.action_names = get_keymap_and_action_names(replay_keymap_name)
+        assert episode_dir.is_dir()
+        self._paths = {}
+        for mode in ['train', 'test', 'imagination']:
+            directory = episode_dir / mode
+            if directory.is_dir():
+                self._paths[mode] = sorted([p for p in directory.iterdir() if 'episode_' in p.stem and p.suffix == '.pt'])
+                print(f'Found {len(self._paths[mode])} {mode} episodes.')
+            else:
+                print(f'No {mode} episodes.')
+
+        self._t, self._episode = None, None
+        self._ep_idx = 0
+        self._mode = 'train'
+        self.load()
+
+    def load(self):
+        self._episode = Episode(**torch.load(self.paths[self._ep_idx]))
+        self._t = 0
+
+    def load_next(self):
+        self._ep_idx = (self._ep_idx + 1) % len(self.paths)
+        self.load()
+
+    def load_previous(self):
+        self._ep_idx = (self._ep_idx - 1) % len(self.paths)
+        self.load()
+
+    def set_mode(self, mode):
+        assert mode in ['train', 'test', 'imagination']
+        if mode in self._paths:
+            self._mode = mode
+            self._ep_idx = 0
+            self.load()
+        else:
+            print(f'No {mode} episodes.')
+
+    def __len__(self):
+        return len(self.ends)
+
+    @property
+    def paths(self):
+        return self._paths[self._mode]
+
+    @property
+    def observations(self):
+        return self._episode.observations
+
+    @property
+    def actions(self):
+        return self._episode.actions
+
+    @property
+    def rewards(self):
+        return self._episode.rewards
+
+    @property
+    def ends(self):
+        return self._episode.ends
+
+    def reset(self):
+        return self.observations[self._t]
+
+    def step(self, action) -> torch.FloatTensor:
+        if action == 1:
+            self._t = (self._t - 1) % len(self)
+        elif action == 2:
+            self._t = (self._t + 1) % len(self)
+        if action == 3:
+            self._t = (self._t - 10) % len(self)
+        elif action == 4:
+            self._t = (self._t + 10) % len(self)
+        elif action == 5:
+            self._t = 0
+        elif action == 6:
+            self.load_previous()
+        elif action == 7:
+            self.load_next()
+        elif action == 8:
+            self.set_mode('train')
+        elif action == 9:
+            self.set_mode('test')
+        elif action == 10:
+            self.set_mode('imagination')
+        act = self.actions[self._t]
+        reward = self.rewards[self._t].item()
+        done = self.ends[self._t].item()
+        info = {
+            'ep_name': f'[{self._mode}] {self.paths[self._ep_idx].stem}',
+            'timestep': self._t,
+            'action': self.action_names[act],
+            'cum_reward': f'{sum(self.rewards[:self._t + 1]):.3f}'
+        }
+        return self.observations[self._t], reward, done, info
+
+    def render(self) -> Image.Image:
+        obs = self.observations[self._t]  # (C, H, W) in [0., 1.]
+        arr = obs.permute(1, 2, 0).numpy().astype(np.uint8)
+        return Image.fromarray(arr)
@@ -0,0 +1,111 @@
+from typing import Tuple, Union
+
+import gym
+import numpy as np
+import pygame
+from PIL import Image
+
+from envs import WorldModelEnv
+from game.keymap import get_keymap_and_action_names
+
+
+class Game:
+    def __init__(self, env: Union[gym.Env, WorldModelEnv], keymap_name: str, size: Tuple[int, int], fps: int, verbose: bool) -> None:
+        self.env = env
+        self.height, self.width = size
+        self.fps = fps
+        self.verbose = verbose
+
+        self.keymap, self.action_names = get_keymap_and_action_names(keymap_name)
+
+        print('Actions:')
+        for key, idx in self.keymap.items():
+            print(f'{pygame.key.name(key)}: {self.action_names[idx]}')
+
+    def run(self) -> None:
+        pygame.init()
+
+        header_height = 100 if self.verbose else 0
+        font_size = 24
+        screen = pygame.display.set_mode((self.width, self.height + header_height))
+        clock = pygame.time.Clock()
+        font = pygame.font.SysFont(None, font_size)
+        header_rect = pygame.Rect(0, 0, self.width, header_height)
+
+        def clear_header():
+            pygame.draw.rect(screen, pygame.Color('black'), header_rect)
+            pygame.draw.rect(screen, pygame.Color('white'), header_rect, 1)
+
+        def draw_text(text, idx_line, idx_column=0):
+            pos = (5 + idx_column * int(self.width // 4), 5 + idx_line * font_size)
+            assert (0 <= pos[0] <= self.width) and (0 <= pos[1] <= header_height)
+            screen.blit(font.render(text, True, pygame.Color('white')), pos)
+
+        def draw_game(image):
+            if isinstance(image, np.ndarray):
+                image = Image.fromarray(image)
+            else:
+                assert isinstance(image, Image.Image)
+            pygame_image = np.array(image.resize((self.width, self.height), resample=Image.NEAREST)).transpose((1, 0, 2))
+            surface = pygame.surfarray.make_surface(pygame_image)
+            screen.blit(surface, (0, header_height))
+
+        if isinstance(self.env, gym.Env):
+            _, info = self.env.reset(return_info=True)
+            draw_game(info['rgb'])
+        else:
+            self.env.reset()
+            draw_game(self.env.render())
+
+        clear_header()
+        pygame.display.flip()
+
+        do_reset, do_wait = False, False
+        should_stop = False
+        while not should_stop:
+
+            action = 0  # noop
+            pygame.event.pump()
+            for event in pygame.event.get():
+                if event.type == pygame.QUIT:
+                    should_stop = True
+                if event.type == pygame.KEYDOWN and event.key in self.keymap.keys():
+                    action = self.keymap[event.key]
+                if event.type == pygame.KEYDOWN and event.key == pygame.K_RETURN:
+                    do_reset = True
+                if event.type == pygame.KEYDOWN and event.key == pygame.K_PERIOD:
+                    do_wait = not do_wait
+            if action == 0:
+                pressed = pygame.key.get_pressed()
+                for key, action in self.keymap.items():
+                    if pressed[key]:
+                        break
+                else:
+                    action = 0
+
+            if do_wait:
+                continue
+
+            _, reward, done, info = self.env.step(action)
+
+            img = info['rgb'] if isinstance(self.env, gym.Env) else self.env.render()
+            draw_game(img)
+
+            if self.verbose:
+                clear_header()
+                draw_text(f'Action: {self.action_names[action]}', idx_line=0)
+                draw_text(f'Reward: {reward if isinstance(reward, float) else reward.item(): .2f}', idx_line=1)
+                draw_text(f'Done: {done}', idx_line=2)
+                if info is not None:
+                    assert isinstance(info, dict)
+                    for i, (k, v) in enumerate(info.items()):
+                        draw_text(f'{k}: {v}', idx_line=i, idx_column=1)
+
+            pygame.display.flip()   # update screen
+            clock.tick(self.fps)    # ensures game maintains the given frame rate
+
+            if do_reset or done:
+                self.env.reset()
+                do_reset = False
+
+        pygame.quit()
@@ -0,0 +1,103 @@
+import gym
+import pygame
+
+
+def get_keymap_and_action_names(name):
+
+    if name == 'empty':
+        return EMPTY_KEYMAP, EMPTY_ACTION_NAMES
+
+    if name == 'episode_replay':
+        return EPISODE_REPLAY_KEYMAP, EPISODE_REPLAY_ACTION_NAMES
+
+    if name == 'atari':
+        return ATARI_KEYMAP, ATARI_ACTION_NAMES
+
+    assert name.startswith('atari/')
+    env_id = name.split('atari/')[1]
+    action_names = [x.lower() for x in gym.make(env_id).get_action_meanings()]
+    keymap = {}
+    for key, value in ATARI_KEYMAP.items():
+        if ATARI_ACTION_NAMES[value] in action_names:
+            keymap[key] = action_names.index(ATARI_ACTION_NAMES[value])
+    return keymap, action_names
+
+
+ATARI_ACTION_NAMES = [
+    'noop',
+    'fire',
+    'up',
+    'right',
+    'left',
+    'down',
+    'upright',
+    'upleft',
+    'downright',
+    'downleft',
+    'upfire',
+    'rightfire',
+    'leftfire',
+    'downfire',
+    'uprightfire',
+    'upleftfire',
+    'downrightfire',
+    'downleftfire',
+]
+
+ATARI_KEYMAP = {
+    pygame.K_SPACE: 1,
+
+    pygame.K_w: 2,
+    pygame.K_d: 3,
+    pygame.K_a: 4,
+    pygame.K_s: 5,
+
+    pygame.K_t: 6,
+    pygame.K_r: 7,
+    pygame.K_g: 8,
+    pygame.K_f: 9,
+
+    pygame.K_UP: 10,
+    pygame.K_RIGHT: 11,
+    pygame.K_LEFT: 12,
+    pygame.K_DOWN: 13,
+
+    pygame.K_u: 14,
+    pygame.K_y: 15,
+    pygame.K_j: 16,
+    pygame.K_h: 17,
+}
+
+EPISODE_REPLAY_ACTION_NAMES = [
+    'noop',
+    'previous',
+    'next',
+    'previous_10',
+    'next_10',
+    'go_to_start',
+    'load_previous',
+    'load_next',
+    'go_to_train_episodes',
+    'go_to_test_episodes',
+    'go_to_imagination_episodes',
+]
+
+EPISODE_REPLAY_KEYMAP = {
+    pygame.K_LEFT: 1,
+    pygame.K_RIGHT: 2,
+    pygame.K_PAGEDOWN: 3,
+    pygame.K_PAGEUP: 4,
+    pygame.K_SPACE: 5,
+    pygame.K_DOWN: 6,
+    pygame.K_UP: 7,
+    pygame.K_t: 8,
+    pygame.K_y: 9,
+    pygame.K_i: 10,
+}
+
+EMPTY_ACTION_NAMES = [
+    'noop',
+]
+
+EMPTY_KEYMAP = {
+}
@@ -0,0 +1,14 @@
+import hydra
+from omegaconf import DictConfig
+
+from trainer import Trainer
+
+
+@hydra.main(config_path="../config", config_name="trainer")
+def main(cfg: DictConfig):
+    trainer = Trainer(cfg)
+    trainer.run()
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,53 @@
+from einops import rearrange
+import numpy as np
+from PIL import Image
+import torch
+
+
+@torch.no_grad()
+def make_reconstructions_from_batch(batch, save_dir, epoch, tokenizer):
+    check_batch(batch)
+
+    original_frames = tensor_to_np_frames(rearrange(batch['observations'], 'b t c h w -> b t h w c'))
+    all = [original_frames]
+
+    rec_frames = generate_reconstructions_with_tokenizer(batch, tokenizer)
+    all.append(rec_frames)
+
+    for i, image in enumerate(map(Image.fromarray, np.concatenate(list(np.concatenate((original_frames, rec_frames), axis=-2)), axis=-3))):
+        image.save(save_dir / f'epoch_{epoch:03d}_t_{i:03d}.png')
+
+    return
+
+
+def check_batch(batch):
+    assert sorted(batch.keys()) == ['actions', 'ends', 'mask_padding', 'observations', 'rewards']
+    b, t, _, _, _ = batch['observations'].shape  # (B, T, C, H, W)
+    assert batch['actions'].shape == batch['rewards'].shape == batch['ends'].shape == batch['mask_padding'].shape == (b, t)
+
+
+def tensor_to_np_frames(inputs):
+    check_float_btw_0_1(inputs)
+    return inputs.mul(255).cpu().numpy().astype(np.uint8)
+
+
+def check_float_btw_0_1(inputs):
+    assert inputs.is_floating_point() and (inputs >= 0).all() and (inputs <= 1).all()
+
+
+@torch.no_grad()
+def generate_reconstructions_with_tokenizer(batch, tokenizer):
+    check_batch(batch)
+    inputs = rearrange(batch['observations'], 'b t c h w -> (b t) c h w')
+    outputs = reconstruct_through_tokenizer(inputs, tokenizer)
+    b, t, _, _, _ = batch['observations'].size()
+    outputs = rearrange(outputs, '(b t) c h w -> b t h w c', b=b, t=t)
+    rec_frames = tensor_to_np_frames(outputs)
+    return rec_frames
+
+
+@torch.no_grad()
+def reconstruct_through_tokenizer(inputs, tokenizer):
+    check_float_btw_0_1(inputs)
+    reconstructions = tokenizer.encode_decode(inputs, should_preprocess=True, should_postprocess=True)
+    return torch.clamp(reconstructions, 0, 1)
@@ -0,0 +1 @@
+from .transformer import Transformer, TransformerConfig
@@ -0,0 +1,166 @@
+from dataclasses import dataclass
+from typing import Any, Optional, Union
+import sys
+
+from einops import rearrange
+import numpy as np
+import torch
+from torch.distributions.categorical import Categorical
+import torch.nn as nn
+import torch.nn.functional as F
+from tqdm import tqdm
+
+from dataset import Batch
+from envs.world_model_env import WorldModelEnv
+from models.tokenizer import Tokenizer
+from models.world_model import WorldModel
+from utils import compute_lambda_returns, LossWithIntermediateLosses
+
+
+@dataclass
+class ActorCriticOutput:
+    logits_actions: torch.FloatTensor
+    means_values: torch.FloatTensor
+
+
+@dataclass
+class ImagineOutput:
+    observations: torch.ByteTensor
+    actions: torch.LongTensor
+    logits_actions: torch.FloatTensor
+    values: torch.FloatTensor
+    rewards: torch.FloatTensor
+    ends: torch.BoolTensor
+
+
+class ActorCritic(nn.Module):
+    def __init__(self, act_vocab_size, use_original_obs: bool = False) -> None:
+        super().__init__()
+        self.use_original_obs = use_original_obs
+        self.conv1 = nn.Conv2d(3, 32, 3, stride=1, padding=1)
+        self.maxp1 = nn.MaxPool2d(2, 2)
+        self.conv2 = nn.Conv2d(32, 32, 3, stride=1, padding=1)
+        self.maxp2 = nn.MaxPool2d(2, 2)
+        self.conv3 = nn.Conv2d(32, 64, 3, stride=1, padding=1)
+        self.maxp3 = nn.MaxPool2d(2, 2)
+        self.conv4 = nn.Conv2d(64, 64, 3, stride=1, padding=1)
+        self.maxp4 = nn.MaxPool2d(2, 2)
+
+        self.lstm_dim = 512
+        self.lstm = nn.LSTMCell(1024, self.lstm_dim)
+        self.hx, self.cx = None, None
+
+        self.critic_linear = nn.Linear(512, 1)
+        self.actor_linear = nn.Linear(512, act_vocab_size)
+
+    def __repr__(self) -> str:
+        return "actor_critic"
+
+    def clear(self) -> None:
+        self.hx, self.cx = None, None
+
+    def reset(self, n: int, burnin_observations: Optional[torch.Tensor] = None, mask_padding: Optional[torch.Tensor] = None) -> None:
+        device = self.conv1.weight.device
+        self.hx = torch.zeros(n, self.lstm_dim, device=device)
+        self.cx = torch.zeros(n, self.lstm_dim, device=device)
+        if burnin_observations is not None:
+            assert burnin_observations.ndim == 5 and burnin_observations.size(0) == n and mask_padding is not None and burnin_observations.shape[:2] == mask_padding.shape
+            for i in range(burnin_observations.size(1)):
+                if mask_padding[:, i].any():
+                    with torch.no_grad():
+                        self(burnin_observations[:, i], mask_padding[:, i])
+
+    def prune(self, mask: np.ndarray) -> None:
+        self.hx = self.hx[mask]
+        self.cx = self.cx[mask]
+
+    def forward(self, inputs: torch.FloatTensor, mask_padding: Optional[torch.BoolTensor] = None) -> ActorCriticOutput:
+        assert inputs.ndim == 4 and inputs.shape[1:] == (3, 64, 64)
+        assert 0 <= inputs.min() <= 1 and 0 <= inputs.max() <= 1
+        assert mask_padding is None or (mask_padding.ndim == 1 and mask_padding.size(0) == inputs.size(0) and mask_padding.any())
+        x = inputs[mask_padding] if mask_padding is not None else inputs
+
+        x = x.mul(2).sub(1)
+        x = F.relu(self.maxp1(self.conv1(x)))
+        x = F.relu(self.maxp2(self.conv2(x)))
+        x = F.relu(self.maxp3(self.conv3(x)))
+        x = F.relu(self.maxp4(self.conv4(x)))
+        x = torch.flatten(x, start_dim=1)
+
+        if mask_padding is None:
+            self.hx, self.cx = self.lstm(x, (self.hx, self.cx))
+        else:
+            self.hx[mask_padding], self.cx[mask_padding] = self.lstm(x, (self.hx[mask_padding], self.cx[mask_padding]))
+
+        logits_actions = rearrange(self.actor_linear(self.hx), 'b a -> b 1 a')
+        means_values = rearrange(self.critic_linear(self.hx), 'b 1 -> b 1 1')
+
+        return ActorCriticOutput(logits_actions, means_values)
+
+    def compute_loss(self, batch: Batch, tokenizer: Tokenizer, world_model: WorldModel, imagine_horizon: int, gamma: float, lambda_: float, entropy_weight: float, **kwargs: Any) -> LossWithIntermediateLosses:
+        assert not self.use_original_obs
+        outputs = self.imagine(batch, tokenizer, world_model, horizon=imagine_horizon)
+
+        with torch.no_grad():
+            lambda_returns = compute_lambda_returns(
+                rewards=outputs.rewards,
+                values=outputs.values,
+                ends=outputs.ends,
+                gamma=gamma,
+                lambda_=lambda_,
+            )[:, :-1]
+
+        values = outputs.values[:, :-1]
+
+        d = Categorical(logits=outputs.logits_actions[:, :-1])
+        log_probs = d.log_prob(outputs.actions[:, :-1])
+        loss_actions = -1 * (log_probs * (lambda_returns - values.detach())).mean()
+        loss_entropy = - entropy_weight * d.entropy().mean()
+        loss_values = F.mse_loss(values, lambda_returns)
+
+        return LossWithIntermediateLosses(loss_actions=loss_actions, loss_values=loss_values, loss_entropy=loss_entropy)
+
+    def imagine(self, batch: Batch, tokenizer: Tokenizer, world_model: WorldModel, horizon: int, show_pbar: bool = False) -> ImagineOutput:
+        assert not self.use_original_obs
+        initial_observations = batch['observations']
+        mask_padding = batch['mask_padding']
+        assert initial_observations.ndim == 5 and initial_observations.shape[2:] == (3, 64, 64)
+        assert mask_padding[:, -1].all()
+        device = initial_observations.device
+        wm_env = WorldModelEnv(tokenizer, world_model, device)
+
+        all_actions = []
+        all_logits_actions = []
+        all_values = []
+        all_rewards = []
+        all_ends = []
+        all_observations = []
+
+        burnin_observations = torch.clamp(tokenizer.encode_decode(initial_observations[:, :-1], should_preprocess=True, should_postprocess=True), 0, 1) if initial_observations.size(1) > 1 else None
+        self.reset(n=initial_observations.size(0), burnin_observations=burnin_observations, mask_padding=mask_padding[:, :-1])
+
+        obs = wm_env.reset_from_initial_observations(initial_observations[:, -1])
+        for k in tqdm(range(horizon), disable=not show_pbar, desc='Imagination', file=sys.stdout):
+
+            all_observations.append(obs)
+
+            outputs_ac = self(obs)
+            action_token = Categorical(logits=outputs_ac.logits_actions).sample()
+            obs, reward, done, _ = wm_env.step(action_token, should_predict_next_obs=(k < horizon - 1))
+
+            all_actions.append(action_token)
+            all_logits_actions.append(outputs_ac.logits_actions)
+            all_values.append(outputs_ac.means_values)
+            all_rewards.append(torch.tensor(reward).reshape(-1, 1))
+            all_ends.append(torch.tensor(done).reshape(-1, 1))
+
+        self.clear()
+
+        return ImagineOutput(
+            observations=torch.stack(all_observations, dim=1).mul(255).byte(),      # (B, T, C, H, W) in [0, 255]
+            actions=torch.cat(all_actions, dim=1),                                  # (B, T)
+            logits_actions=torch.cat(all_logits_actions, dim=1),                    # (B, T, #actions)
+            values=rearrange(torch.cat(all_values, dim=1), 'b t 1 -> b t'),         # (B, T)
+            rewards=torch.cat(all_rewards, dim=1).to(device),                       # (B, T)
+            ends=torch.cat(all_ends, dim=1).to(device),                             # (B, T)
+        )
@@ -0,0 +1,106 @@
+from typing import Tuple
+
+import numpy as np
+import torch
+
+
+class Cache:
+    def __init__(self, num_samples: int, num_heads: int, max_tokens: int, embed_dim: int, device: torch.device) -> None:
+        assert embed_dim % num_heads == 0
+        self._n, self._cache, self._size = num_samples, None, None
+        self._reset = lambda n: torch.empty(n, num_heads, max_tokens, embed_dim // num_heads, device=device)  # (B, nh, T, hs)
+        self.reset()
+
+    @property
+    def shape(self) -> Tuple[int, int, int, int]:
+        n, num_heads, _, head_dim = self._cache.shape
+        return n, num_heads, self._size, head_dim
+
+    def reset(self) -> None:
+        self._cache = self._reset(self._n)
+        self._size = 0
+
+    def prune(self, mask: np.ndarray) -> None:
+        assert mask.ndim == 1 and mask.shape[0] == self.shape[0]
+        self._cache = self._cache[mask]
+        self._n = self._cache.shape[0]
+
+    def get(self) -> torch.Tensor:
+        return self._cache[:, :, :self._size, :]
+
+    def update(self, x: torch.Tensor) -> None:
+        assert (x.ndim == self._cache.ndim) and all([x.size(i) == self._cache.size(i) for i in (0, 1, 3)])
+        assert self._size + x.size(2) <= self._cache.shape[2]
+        self._cache = AssignWithoutInplaceCheck.apply(self._cache, x, 2, self._size, self._size + x.size(2))
+        self._size += x.size(2)
+
+
+class KVCache:
+    def __init__(self, n: int, num_heads: int, max_tokens: int, embed_dim: int, device: torch.device) -> None:
+        self._k_cache = Cache(n, num_heads, max_tokens, embed_dim, device)
+        self._v_cache = Cache(n, num_heads, max_tokens, embed_dim, device)
+
+    @property
+    def shape(self) -> Tuple[int, int, int, int]:
+        return self._k_cache.shape
+
+    def reset(self) -> None:
+        self._k_cache.reset()
+        self._v_cache.reset()
+
+    def prune(self, mask: np.ndarray) -> None:
+        self._k_cache.prune(mask)
+        self._v_cache.prune(mask)
+
+    def get(self) -> Tuple[torch.Tensor, torch.Tensor]:
+        return self._k_cache.get(), self._v_cache.get()
+
+    def update(self, k: torch.Tensor, v: torch.Tensor):
+        self._k_cache.update(k)
+        self._v_cache.update(v)
+
+
+class KeysValues:
+    def __init__(self, n: int, num_heads: int, max_tokens: int, embed_dim: int, num_layers: int, device: torch.device) -> None:
+        self._keys_values = tuple([KVCache(n, num_heads, max_tokens, embed_dim, device) for _ in range(num_layers)])
+
+    def __getitem__(self, key: int) -> KVCache:
+        return self._keys_values[key]
+
+    def __len__(self):
+        return len(self._keys_values)
+
+    @property
+    def size(self):
+        return self._keys_values[0].shape[2]
+
+    def reset(self) -> None:
+        for kv_cache in self._keys_values:
+            kv_cache.reset()
+
+    def prune(self, mask: np.ndarray) -> None:
+        for kv_cache in self._keys_values:
+            kv_cache.prune(mask)
+
+
+class AssignWithoutInplaceCheck(torch.autograd.Function):
+    """
+    Inspired from : https://discuss.pytorch.org/t/disable-in-place-correctness-version-check-any-other-workaround/90738/4
+    Warning : do not use it to overwrite a slice twice.
+    """
+
+    @staticmethod
+    def get_slice(dim: int, start: int, stop: int) -> Tuple[slice]:
+        return tuple([slice(None), ] * dim + [slice(start, stop)])
+
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, value: torch.Tensor, dim: int, start: int, stop: int) -> torch.Tensor:
+        ctx.dim = dim
+        ctx.start = start
+        ctx.stop = stop
+        input.data[AssignWithoutInplaceCheck.get_slice(dim, start, stop)] = value
+        return input
+
+    @staticmethod
+    def backward(ctx, grad_out: torch.Tensor) -> Tuple[torch.Tensor]:
+        return grad_out, grad_out[AssignWithoutInplaceCheck.get_slice(ctx.dim, ctx.start, ctx.stop)], None, None, None
@@ -0,0 +1,54 @@
+import math
+from typing import List
+
+import torch
+import torch.nn as nn
+
+
+class Slicer(nn.Module):
+    def __init__(self, max_blocks: int, block_mask: torch.Tensor) -> None:
+        super().__init__()
+        self.block_size = block_mask.size(0)
+        self.num_kept_tokens = block_mask.sum().long().item()
+        kept_indices = torch.where(block_mask)[0].repeat(max_blocks)
+        offsets = torch.arange(max_blocks).repeat_interleave(self.num_kept_tokens)
+        self.register_buffer('indices', kept_indices + block_mask.size(0) * offsets)
+
+    def compute_slice(self, num_steps: int, prev_steps: int = 0) -> torch.Tensor:
+        total_steps = num_steps + prev_steps
+        num_blocks = math.ceil(total_steps / self.block_size)
+        indices = self.indices[:num_blocks * self.num_kept_tokens]
+        return indices[torch.logical_and(prev_steps <= indices, indices < total_steps)] - prev_steps
+
+    def forward(self, *args, **kwargs):
+        raise NotImplementedError
+
+
+class Head(Slicer):
+    def __init__(self, max_blocks: int, block_mask: torch.Tensor, head_module: nn.Module) -> None:
+        super().__init__(max_blocks, block_mask)
+        assert isinstance(head_module, nn.Module)
+        self.head_module = head_module
+
+    def forward(self, x: torch.Tensor, num_steps: int, prev_steps: int) -> torch.Tensor:
+        x_sliced = x[:, self.compute_slice(num_steps, prev_steps)]  # x is (B, T, E)
+        return self.head_module(x_sliced)
+
+
+class Embedder(nn.Module):
+    def __init__(self, max_blocks: int, block_masks: List[torch.Tensor], embedding_tables: List[nn.Embedding]) -> None:
+        super().__init__()
+        assert len(block_masks) == len(embedding_tables)
+        assert (sum(block_masks) == 1).all()  # block mask are a partition of a block
+        self.embedding_dim = embedding_tables[0].embedding_dim
+        assert all([e.embedding_dim == self.embedding_dim for e in embedding_tables])
+        self.embedding_tables = embedding_tables
+        self.slicers = [Slicer(max_blocks, block_mask) for block_mask in block_masks]
+
+    def forward(self, tokens: torch.Tensor, num_steps: int, prev_steps: int) -> torch.Tensor:
+        assert tokens.ndim == 2  # x is (B, T)
+        output = torch.zeros(*tokens.size(), self.embedding_dim, device=tokens.device)
+        for slicer, emb in zip(self.slicers, self.embedding_tables):
+            s = slicer.compute_slice(num_steps, prev_steps)
+            output[:, s] = emb(tokens[:, s])
+        return output
@@ -0,0 +1,2 @@
+from .nets import Encoder, Decoder, EncoderDecoderConfig
+from .tokenizer import Tokenizer, TokenizerEncoderOutput
@@ -0,0 +1,167 @@
+"""
+Credits to https://github.com/CompVis/taming-transformers
+"""
+
+from collections import namedtuple
+import hashlib
+import os
+from pathlib import Path
+import requests
+
+import torch
+import torch.nn as nn
+from torchvision import models
+from tqdm import tqdm
+
+
+class LPIPS(nn.Module):
+    # Learned perceptual metric
+    def __init__(self, use_dropout: bool = True):
+        super().__init__()
+        self.scaling_layer = ScalingLayer()
+        self.chns = [64, 128, 256, 512, 512]  # vg16 features
+        self.net = vgg16(pretrained=True, requires_grad=False)
+        self.lin0 = NetLinLayer(self.chns[0], use_dropout=use_dropout)
+        self.lin1 = NetLinLayer(self.chns[1], use_dropout=use_dropout)
+        self.lin2 = NetLinLayer(self.chns[2], use_dropout=use_dropout)
+        self.lin3 = NetLinLayer(self.chns[3], use_dropout=use_dropout)
+        self.lin4 = NetLinLayer(self.chns[4], use_dropout=use_dropout)
+        self.load_from_pretrained()
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def load_from_pretrained(self) -> None:
+        ckpt = get_ckpt_path(name="vgg_lpips", root=Path.home() / ".cache/iris/tokenizer_pretrained_vgg")  # Download VGG if necessary
+        self.load_state_dict(torch.load(ckpt, map_location=torch.device("cpu")), strict=False)
+
+    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        in0_input, in1_input = (self.scaling_layer(input), self.scaling_layer(target))
+        outs0, outs1 = self.net(in0_input), self.net(in1_input)
+        feats0, feats1, diffs = {}, {}, {}
+        lins = [self.lin0, self.lin1, self.lin2, self.lin3, self.lin4]
+        for kk in range(len(self.chns)):
+            feats0[kk], feats1[kk] = normalize_tensor(outs0[kk]), normalize_tensor(outs1[kk])
+            diffs[kk] = (feats0[kk] - feats1[kk]) ** 2
+
+        res = [spatial_average(lins[kk].model(diffs[kk]), keepdim=True) for kk in range(len(self.chns))]
+        val = res[0]
+        for i in range(1, len(self.chns)):
+            val += res[i]
+        return val
+
+
+class ScalingLayer(nn.Module):
+    def __init__(self) -> None:
+        super(ScalingLayer, self).__init__()
+        self.register_buffer('shift', torch.Tensor([-.030, -.088, -.188])[None, :, None, None])
+        self.register_buffer('scale', torch.Tensor([.458, .448, .450])[None, :, None, None])
+
+    def forward(self, inp: torch.Tensor) -> torch.Tensor:
+        return (inp - self.shift) / self.scale
+
+
+class NetLinLayer(nn.Module):
+    """ A single linear layer which does a 1x1 conv """
+    def __init__(self, chn_in: int, chn_out: int = 1, use_dropout: bool = False) -> None:
+        super(NetLinLayer, self).__init__()
+        layers = [nn.Dropout(), ] if (use_dropout) else []
+        layers += [nn.Conv2d(chn_in, chn_out, 1, stride=1, padding=0, bias=False), ]
+        self.model = nn.Sequential(*layers)
+
+
+class vgg16(torch.nn.Module):
+    def __init__(self, requires_grad: bool = False, pretrained: bool = True) -> None:
+        super(vgg16, self).__init__()
+        vgg_pretrained_features = models.vgg16(pretrained=pretrained).features
+        self.slice1 = torch.nn.Sequential()
+        self.slice2 = torch.nn.Sequential()
+        self.slice3 = torch.nn.Sequential()
+        self.slice4 = torch.nn.Sequential()
+        self.slice5 = torch.nn.Sequential()
+        self.N_slices = 5
+        for x in range(4):
+            self.slice1.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(4, 9):
+            self.slice2.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(9, 16):
+            self.slice3.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(16, 23):
+            self.slice4.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(23, 30):
+            self.slice5.add_module(str(x), vgg_pretrained_features[x])
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, X: torch.Tensor) -> torch.Tensor:
+        h = self.slice1(X)
+        h_relu1_2 = h
+        h = self.slice2(h)
+        h_relu2_2 = h
+        h = self.slice3(h)
+        h_relu3_3 = h
+        h = self.slice4(h)
+        h_relu4_3 = h
+        h = self.slice5(h)
+        h_relu5_3 = h
+        vgg_outputs = namedtuple("VggOutputs", ['relu1_2', 'relu2_2', 'relu3_3', 'relu4_3', 'relu5_3'])
+        out = vgg_outputs(h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3, h_relu5_3)
+        return out
+
+
+def normalize_tensor(x: torch.Tensor, eps: float = 1e-10) -> torch.Tensor:
+    norm_factor = torch.sqrt(torch.sum(x ** 2, dim=1, keepdim=True))
+    return x / (norm_factor + eps)
+
+
+def spatial_average(x: torch.Tensor, keepdim: bool = True) -> torch.Tensor:
+    return x.mean([2, 3], keepdim=keepdim)
+
+
+# ********************************************************************
+# *************** Utilities to download pretrained vgg ***************
+# ********************************************************************
+
+
+URL_MAP = {
+    "vgg_lpips": "https://heibox.uni-heidelberg.de/f/607503859c864bc1b30b/?dl=1"
+}
+
+
+CKPT_MAP = {
+    "vgg_lpips": "vgg.pth"
+}
+
+
+MD5_MAP = {
+    "vgg_lpips": "d507d7349b931f0638a25a48a722f98a"
+}
+
+
+def download(url: str, local_path: str, chunk_size: int = 1024) -> None:
+    os.makedirs(os.path.split(local_path)[0], exist_ok=True)
+    with requests.get(url, stream=True) as r:
+        total_size = int(r.headers.get("content-length", 0))
+        with tqdm(total=total_size, unit="B", unit_scale=True) as pbar:
+            with open(local_path, "wb") as f:
+                for data in r.iter_content(chunk_size=chunk_size):
+                    if data:
+                        f.write(data)
+                        pbar.update(chunk_size)
+
+
+def md5_hash(path: str) -> str:
+    with open(path, "rb") as f:
+        content = f.read()
+    return hashlib.md5(content).hexdigest()
+
+
+def get_ckpt_path(name: str, root: str, check: bool = False) -> str:
+    assert name in URL_MAP
+    path = os.path.join(root, CKPT_MAP[name])
+    if not os.path.exists(path) or (check and not md5_hash(path) == MD5_MAP[name]):
+        print("Downloading {} model from {} to {}".format(name, URL_MAP[name], path))
+        download(URL_MAP[name], path)
+        md5 = md5_hash(path)
+        assert md5 == MD5_MAP[name], md5
+    return path
@@ -0,0 +1,362 @@
+"""
+Credits to https://github.com/CompVis/taming-transformers
+"""
+
+from dataclasses import dataclass
+from typing import List
+
+import torch
+import torch.nn as nn
+
+
+@dataclass
+class EncoderDecoderConfig:
+    resolution: int
+    in_channels: int
+    z_channels: int
+    ch: int
+    ch_mult: List[int]
+    num_res_blocks: int
+    attn_resolutions: List[int]
+    out_ch: int
+    dropout: float
+
+
+class Encoder(nn.Module):
+    def __init__(self, config: EncoderDecoderConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.num_resolutions = len(config.ch_mult)
+        temb_ch = 0  # timestep embedding #channels
+
+        # downsampling
+        self.conv_in = torch.nn.Conv2d(config.in_channels,
+                                       config.ch,
+                                       kernel_size=3,
+                                       stride=1,
+                                       padding=1)
+
+        curr_res = config.resolution
+        in_ch_mult = (1,) + tuple(config.ch_mult)
+        self.down = nn.ModuleList()
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = config.ch * in_ch_mult[i_level]
+            block_out = config.ch * config.ch_mult[i_level]
+            for i_block in range(self.config.num_res_blocks):
+                block.append(ResnetBlock(in_channels=block_in,
+                                         out_channels=block_out,
+                                         temb_channels=temb_ch,
+                                         dropout=config.dropout))
+                block_in = block_out
+                if curr_res in config.attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level != self.num_resolutions - 1:
+                down.downsample = Downsample(block_in, with_conv=True)
+                curr_res = curr_res // 2
+            self.down.append(down)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=temb_ch,
+                                       dropout=config.dropout)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=temb_ch,
+                                       dropout=config.dropout)
+
+        # end
+        self.norm_out = Normalize(block_in)
+        self.conv_out = torch.nn.Conv2d(block_in,
+                                        config.z_channels,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+
+        temb = None  # timestep embedding
+
+        # downsampling
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.config.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1], temb)
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions - 1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        # middle
+        h = hs[-1]
+        h = self.mid.block_1(h, temb)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h, temb)
+
+        # end
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        return h
+
+
+class Decoder(nn.Module):
+    def __init__(self, config: EncoderDecoderConfig) -> None:
+        super().__init__()
+        self.config = config
+        temb_ch = 0
+        self.num_resolutions = len(config.ch_mult)
+
+        # compute in_ch_mult, block_in and curr_res at lowest res
+        in_ch_mult = (1,) + tuple(config.ch_mult)
+        block_in = config.ch * config.ch_mult[self.num_resolutions - 1]
+        curr_res = config.resolution // 2 ** (self.num_resolutions - 1)
+        print(f"Tokenizer : shape of latent is {config.z_channels, curr_res, curr_res}.")
+
+        # z to block_in
+        self.conv_in = torch.nn.Conv2d(config.z_channels,
+                                       block_in,
+                                       kernel_size=3,
+                                       stride=1,
+                                       padding=1)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=temb_ch,
+                                       dropout=config.dropout)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=temb_ch,
+                                       dropout=config.dropout)
+
+        # upsampling
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = config.ch * config.ch_mult[i_level]
+            for i_block in range(config.num_res_blocks + 1):
+                block.append(ResnetBlock(in_channels=block_in,
+                                         out_channels=block_out,
+                                         temb_channels=temb_ch,
+                                         dropout=config.dropout))
+                block_in = block_out
+                if curr_res in config.attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level != 0:
+                up.upsample = Upsample(block_in, with_conv=True)
+                curr_res = curr_res * 2
+            self.up.insert(0, up)  # prepend to get consistent order
+
+        # end
+        self.norm_out = Normalize(block_in)
+        self.conv_out = torch.nn.Conv2d(block_in,
+                                        config.out_ch,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+    def forward(self, z: torch.Tensor) -> torch.Tensor:
+        temb = None  # timestep embedding
+
+        # z to block_in
+        h = self.conv_in(z)
+
+        # middle
+        h = self.mid.block_1(h, temb)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h, temb)
+
+        # upsampling
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.config.num_res_blocks + 1):
+                h = self.up[i_level].block[i_block](h, temb)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+
+        # end
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        return h
+
+
+def nonlinearity(x: torch.Tensor) -> torch.Tensor:
+    # swish
+    return x * torch.sigmoid(x)
+
+
+def Normalize(in_channels: int) -> nn.Module:
+    return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+
+
+class Upsample(nn.Module):
+    def __init__(self, in_channels: int, with_conv: bool) -> None:
+        super().__init__()
+        self.with_conv = with_conv
+        if self.with_conv:
+            self.conv = torch.nn.Conv2d(in_channels,
+                                        in_channels,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
+        if self.with_conv:
+            x = self.conv(x)
+        return x
+
+
+class Downsample(nn.Module):
+    def __init__(self, in_channels: int, with_conv: bool) -> None:
+        super().__init__()
+        self.with_conv = with_conv
+        if self.with_conv:
+            # no asymmetric padding in torch conv, must do it ourselves
+            self.conv = torch.nn.Conv2d(in_channels,
+                                        in_channels,
+                                        kernel_size=3,
+                                        stride=2,
+                                        padding=0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.with_conv:
+            pad = (0, 1, 0, 1)
+            x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
+            x = self.conv(x)
+        else:
+            x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
+        return x
+
+
+class ResnetBlock(nn.Module):
+    def __init__(self, *, in_channels: int, out_channels: int = None, conv_shortcut: bool = False,
+                 dropout: float, temb_channels: int = 512) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.use_conv_shortcut = conv_shortcut
+
+        self.norm1 = Normalize(in_channels)
+        self.conv1 = torch.nn.Conv2d(in_channels,
+                                     out_channels,
+                                     kernel_size=3,
+                                     stride=1,
+                                     padding=1)
+        if temb_channels > 0:
+            self.temb_proj = torch.nn.Linear(temb_channels,
+                                             out_channels)
+        self.norm2 = Normalize(out_channels)
+        self.dropout = torch.nn.Dropout(dropout)
+        self.conv2 = torch.nn.Conv2d(out_channels,
+                                     out_channels,
+                                     kernel_size=3,
+                                     stride=1,
+                                     padding=1)
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                self.conv_shortcut = torch.nn.Conv2d(in_channels,
+                                                     out_channels,
+                                                     kernel_size=3,
+                                                     stride=1,
+                                                     padding=1)
+            else:
+                self.nin_shortcut = torch.nn.Conv2d(in_channels,
+                                                    out_channels,
+                                                    kernel_size=1,
+                                                    stride=1,
+                                                    padding=0)
+
+    def forward(self, x: torch.Tensor, temb: torch.Tensor) -> torch.Tensor:
+        h = x
+        h = self.norm1(h)
+        h = nonlinearity(h)
+        h = self.conv1(h)
+
+        if temb is not None:
+            h = h + self.temb_proj(nonlinearity(temb))[:, :, None, None]
+
+        h = self.norm2(h)
+        h = nonlinearity(h)
+        h = self.dropout(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                x = self.conv_shortcut(x)
+            else:
+                x = self.nin_shortcut(x)
+
+        return x + h
+
+
+class AttnBlock(nn.Module):
+    def __init__(self, in_channels: int) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = Normalize(in_channels)
+        self.q = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.k = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.v = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.proj_out = torch.nn.Conv2d(in_channels,
+                                        in_channels,
+                                        kernel_size=1,
+                                        stride=1,
+                                        padding=0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h_ = x
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        # compute attention
+        b, c, h, w = q.shape
+        q = q.reshape(b, c, h * w)
+        q = q.permute(0, 2, 1)      # b,hw,c
+        k = k.reshape(b, c, h * w)  # b,c,hw
+        w_ = torch.bmm(q, k)        # b,hw,hw    w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
+        w_ = w_ * (int(c) ** (-0.5))
+        w_ = torch.nn.functional.softmax(w_, dim=2)
+
+        # attend to values
+        v = v.reshape(b, c, h * w)
+        w_ = w_.permute(0, 2, 1)   # b,hw,hw (first hw of k, second of q)
+        h_ = torch.bmm(v, w_)     # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
+        h_ = h_.reshape(b, c, h, w)
+
+        h_ = self.proj_out(h_)
+
+        return x + h_
@@ -0,0 +1,104 @@
+"""
+Credits to https://github.com/CompVis/taming-transformers
+"""
+
+from dataclasses import dataclass
+from typing import Any, Tuple
+
+from einops import rearrange
+import torch
+import torch.nn as nn
+
+from dataset import Batch
+from .lpips import LPIPS
+from .nets import Encoder, Decoder
+from utils import LossWithIntermediateLosses
+
+
+@dataclass
+class TokenizerEncoderOutput:
+    z: torch.FloatTensor
+    z_quantized: torch.FloatTensor
+    tokens: torch.LongTensor
+
+
+class Tokenizer(nn.Module):
+    def __init__(self, vocab_size: int, embed_dim: int, encoder: Encoder, decoder: Decoder, with_lpips: bool = True) -> None:
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.encoder = encoder
+        self.pre_quant_conv = torch.nn.Conv2d(encoder.config.z_channels, embed_dim, 1)
+        self.embedding = nn.Embedding(vocab_size, embed_dim)
+        self.post_quant_conv = torch.nn.Conv2d(embed_dim, decoder.config.z_channels, 1)
+        self.decoder = decoder
+        self.embedding.weight.data.uniform_(-1.0 / vocab_size, 1.0 / vocab_size)
+        self.lpips = LPIPS().eval() if with_lpips else None
+
+    def __repr__(self) -> str:
+        return "tokenizer"
+
+    def forward(self, x: torch.Tensor, should_preprocess: bool = False, should_postprocess: bool = False) -> Tuple[torch.Tensor]:
+        outputs = self.encode(x, should_preprocess)
+        decoder_input = outputs.z + (outputs.z_quantized - outputs.z).detach()
+        reconstructions = self.decode(decoder_input, should_postprocess)
+        return outputs.z, outputs.z_quantized, reconstructions
+
+    def compute_loss(self, batch: Batch, **kwargs: Any) -> LossWithIntermediateLosses:
+        assert self.lpips is not None
+        observations = self.preprocess_input(rearrange(batch['observations'], 'b t c h w -> (b t) c h w'))
+        z, z_quantized, reconstructions = self(observations, should_preprocess=False, should_postprocess=False)
+
+        # Codebook loss. Notes:
+        # - beta position is different from taming and identical to original VQVAE paper
+        # - VQVAE uses 0.25 by default
+        beta = 1.0
+        commitment_loss = (z.detach() - z_quantized).pow(2).mean() + beta * (z - z_quantized.detach()).pow(2).mean()
+
+        reconstruction_loss = torch.abs(observations - reconstructions).mean()
+        perceptual_loss = torch.mean(self.lpips(observations, reconstructions))
+
+        return LossWithIntermediateLosses(commitment_loss=commitment_loss, reconstruction_loss=reconstruction_loss, perceptual_loss=perceptual_loss)
+
+    def encode(self, x: torch.Tensor, should_preprocess: bool = False) -> TokenizerEncoderOutput:
+        if should_preprocess:
+            x = self.preprocess_input(x)
+        shape = x.shape  # (..., C, H, W)
+        x = x.view(-1, *shape[-3:])
+        z = self.encoder(x)
+        z = self.pre_quant_conv(z)
+        b, e, h, w = z.shape
+        z_flattened = rearrange(z, 'b e h w -> (b h w) e')
+        dist_to_embeddings = torch.sum(z_flattened ** 2, dim=1, keepdim=True) + torch.sum(self.embedding.weight**2, dim=1) - 2 * torch.matmul(z_flattened, self.embedding.weight.t())
+
+        tokens = dist_to_embeddings.argmin(dim=-1)
+        z_q = rearrange(self.embedding(tokens), '(b h w) e -> b e h w', b=b, e=e, h=h, w=w).contiguous()
+
+        # Reshape to original
+        z = z.reshape(*shape[:-3], *z.shape[1:])
+        z_q = z_q.reshape(*shape[:-3], *z_q.shape[1:])
+        tokens = tokens.reshape(*shape[:-3], -1)
+
+        return TokenizerEncoderOutput(z, z_q, tokens)
+
+    def decode(self, z_q: torch.Tensor, should_postprocess: bool = False) -> torch.Tensor:
+        shape = z_q.shape  # (..., E, h, w)
+        z_q = z_q.view(-1, *shape[-3:])
+        z_q = self.post_quant_conv(z_q)
+        rec = self.decoder(z_q)
+        rec = rec.reshape(*shape[:-3], *rec.shape[1:])
+        if should_postprocess:
+            rec = self.postprocess_output(rec)
+        return rec
+
+    @torch.no_grad()
+    def encode_decode(self, x: torch.Tensor, should_preprocess: bool = False, should_postprocess: bool = False) -> torch.Tensor:
+        z_q = self.encode(x, should_preprocess).z_quantized
+        return self.decode(z_q, should_postprocess)
+
+    def preprocess_input(self, x: torch.Tensor) -> torch.Tensor:
+        """x is supposed to be channels first and in [0, 1]"""
+        return x.mul(2).sub(1)
+
+    def postprocess_output(self, y: torch.Tensor) -> torch.Tensor:
+        """y is supposed to be channels first and in [-1, 1]"""
+        return y.add(1).div(2)
@@ -0,0 +1,120 @@
+"""
+Credits to https://github.com/karpathy/minGPT
+"""
+
+from dataclasses import dataclass
+import math
+from typing import Optional
+
+from einops import rearrange
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+from .kv_caching import KeysValues, KVCache
+
+
+@dataclass
+class TransformerConfig:
+    tokens_per_block: int
+    max_blocks: int
+    attention: str
+
+    num_layers: int
+    num_heads: int
+    embed_dim: int
+
+    embed_pdrop: float
+    resid_pdrop: float
+    attn_pdrop: float
+
+    @property
+    def max_tokens(self):
+        return self.tokens_per_block * self.max_blocks
+
+
+class Transformer(nn.Module):
+    def __init__(self, config: TransformerConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.drop = nn.Dropout(config.embed_pdrop)
+        self.blocks = nn.ModuleList([Block(config) for _ in range(config.num_layers)])
+        self.ln_f = nn.LayerNorm(config.embed_dim)
+
+    def generate_empty_keys_values(self, n: int, max_tokens: int) -> KeysValues:
+        device = self.ln_f.weight.device  # Assumption that all submodules are on the same device
+        return KeysValues(n, self.config.num_heads, max_tokens, self.config.embed_dim, self.config.num_layers, device)
+
+    def forward(self, sequences: torch.Tensor, past_keys_values: Optional[KeysValues] = None) -> torch.Tensor:
+        assert past_keys_values is None or len(past_keys_values) == len(self.blocks)
+        x = self.drop(sequences)
+        for i, block in enumerate(self.blocks):
+            x = block(x, None if past_keys_values is None else past_keys_values[i])
+
+        x = self.ln_f(x)
+        return x
+
+
+class Block(nn.Module):
+    def __init__(self, config: TransformerConfig) -> None:
+        super().__init__()
+        self.ln1 = nn.LayerNorm(config.embed_dim)
+        self.ln2 = nn.LayerNorm(config.embed_dim)
+        self.attn = SelfAttention(config)
+        self.mlp = nn.Sequential(
+            nn.Linear(config.embed_dim, 4 * config.embed_dim),
+            nn.GELU(),
+            nn.Linear(4 * config.embed_dim, config.embed_dim),
+            nn.Dropout(config.resid_pdrop),
+        )
+
+    def forward(self, x: torch.Tensor, past_keys_values: Optional[KeysValues] = None) -> torch.Tensor:
+        x_attn = self.attn(self.ln1(x), past_keys_values)
+        x = x + x_attn
+        x = x + self.mlp(self.ln2(x))
+        return x
+
+
+class SelfAttention(nn.Module):
+    def __init__(self, config: TransformerConfig) -> None:
+        super().__init__()
+        assert config.embed_dim % config.num_heads == 0
+        assert config.attention in ('causal', 'block_causal')
+        self.num_heads = config.num_heads
+        self.key = nn.Linear(config.embed_dim, config.embed_dim)
+        self.query = nn.Linear(config.embed_dim, config.embed_dim)
+        self.value = nn.Linear(config.embed_dim, config.embed_dim)
+        self.attn_drop = nn.Dropout(config.attn_pdrop)
+        self.resid_drop = nn.Dropout(config.resid_pdrop)
+        self.proj = nn.Linear(config.embed_dim, config.embed_dim)
+
+        causal_mask = torch.tril(torch.ones(config.max_tokens, config.max_tokens))
+        block_causal_mask = torch.max(causal_mask, torch.block_diag(*[torch.ones(config.tokens_per_block, config.tokens_per_block) for _ in range(config.max_blocks)]))
+        self.register_buffer('mask', causal_mask if config.attention == 'causal' else block_causal_mask)
+
+    def forward(self, x: torch.Tensor, kv_cache: Optional[KVCache] = None) -> torch.Tensor:
+        B, T, C = x.size()
+        if kv_cache is not None:
+            b, nh, L, c = kv_cache.shape
+            assert nh == self.num_heads and b == B and c * nh == C
+        else:
+            L = 0
+
+        q = self.query(x).view(B, T, self.num_heads, C // self.num_heads).transpose(1, 2)   # (B, nh, T, hs)
+        k = self.key(x).view(B, T, self.num_heads, C // self.num_heads).transpose(1, 2)     # (B, nh, T, hs)
+        v = self.value(x).view(B, T, self.num_heads, C // self.num_heads).transpose(1, 2)   # (B, nh, T, hs)
+
+        if kv_cache is not None:
+            kv_cache.update(k, v)
+            k, v = kv_cache.get()
+
+        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
+        att = att.masked_fill(self.mask[L:L + T, :L + T] == 0, float('-inf'))
+        att = F.softmax(att, dim=-1)
+        att = self.attn_drop(att)
+        y = att @ v
+        y = rearrange(y, 'b h t e -> b t (h e)')
+
+        y = self.resid_drop(self.proj(y))
+
+        return y
@@ -0,0 +1,122 @@
+from dataclasses import dataclass
+from typing import Any, Optional, Tuple
+
+from einops import rearrange
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from dataset import Batch
+from .kv_caching import KeysValues
+from .slicer import Embedder, Head
+from .tokenizer import Tokenizer
+from .transformer import Transformer, TransformerConfig
+from utils import init_weights, LossWithIntermediateLosses
+
+
+@dataclass
+class WorldModelOutput:
+    output_sequence: torch.FloatTensor
+    logits_observations: torch.FloatTensor
+    logits_rewards: torch.FloatTensor
+    logits_ends: torch.FloatTensor
+
+
+class WorldModel(nn.Module):
+    def __init__(self, obs_vocab_size: int, act_vocab_size: int, config: TransformerConfig) -> None:
+        super().__init__()
+        self.obs_vocab_size, self.act_vocab_size = obs_vocab_size, act_vocab_size
+        self.config = config
+        self.transformer = Transformer(config)
+
+        all_but_last_obs_tokens_pattern = torch.ones(config.tokens_per_block)
+        all_but_last_obs_tokens_pattern[-2] = 0
+        act_tokens_pattern = torch.zeros(self.config.tokens_per_block)
+        act_tokens_pattern[-1] = 1
+        obs_tokens_pattern = 1 - act_tokens_pattern
+
+        self.pos_emb = nn.Embedding(config.max_tokens, config.embed_dim)
+
+        self.embedder = Embedder(
+            max_blocks=config.max_blocks,
+            block_masks=[act_tokens_pattern, obs_tokens_pattern],
+            embedding_tables=nn.ModuleList([nn.Embedding(act_vocab_size, config.embed_dim), nn.Embedding(obs_vocab_size, config.embed_dim)])
+        )
+
+        self.head_observations = Head(
+            max_blocks=config.max_blocks,
+            block_mask=all_but_last_obs_tokens_pattern,
+            head_module=nn.Sequential(
+                nn.Linear(config.embed_dim, config.embed_dim),
+                nn.ReLU(),
+                nn.Linear(config.embed_dim, obs_vocab_size)
+            )
+        )
+
+        self.head_rewards = Head(
+            max_blocks=config.max_blocks,
+            block_mask=act_tokens_pattern,
+            head_module=nn.Sequential(
+                nn.Linear(config.embed_dim, config.embed_dim),
+                nn.ReLU(),
+                nn.Linear(config.embed_dim, 3)
+            )
+        )
+
+        self.head_ends = Head(
+            max_blocks=config.max_blocks,
+            block_mask=act_tokens_pattern,
+            head_module=nn.Sequential(
+                nn.Linear(config.embed_dim, config.embed_dim),
+                nn.ReLU(),
+                nn.Linear(config.embed_dim, 2)
+            )
+        )
+
+        self.apply(init_weights)
+
+    def __repr__(self) -> str:
+        return "world_model"
+
+    def forward(self, tokens: torch.LongTensor, past_keys_values: Optional[KeysValues] = None) -> WorldModelOutput:
+
+        num_steps = tokens.size(1)  # (B, T)
+        assert num_steps <= self.config.max_tokens
+        prev_steps = 0 if past_keys_values is None else past_keys_values.size
+
+        sequences = self.embedder(tokens, num_steps, prev_steps) + self.pos_emb(prev_steps + torch.arange(num_steps, device=tokens.device))
+
+        x = self.transformer(sequences, past_keys_values)
+
+        logits_observations = self.head_observations(x, num_steps=num_steps, prev_steps=prev_steps)
+        logits_rewards = self.head_rewards(x, num_steps=num_steps, prev_steps=prev_steps)
+        logits_ends = self.head_ends(x, num_steps=num_steps, prev_steps=prev_steps)
+
+        return WorldModelOutput(x, logits_observations, logits_rewards, logits_ends)
+
+    def compute_loss(self, batch: Batch, tokenizer: Tokenizer, **kwargs: Any) -> LossWithIntermediateLosses:
+
+        with torch.no_grad():
+            obs_tokens = tokenizer.encode(batch['observations'], should_preprocess=True).tokens  # (BL, K)
+
+        act_tokens = rearrange(batch['actions'], 'b l -> b l 1')
+        tokens = rearrange(torch.cat((obs_tokens, act_tokens), dim=2), 'b l k1 -> b (l k1)')  # (B, L(K+1))
+
+        outputs = self(tokens)
+
+        labels_observations, labels_rewards, labels_ends = self.compute_labels_world_model(obs_tokens, batch['rewards'], batch['ends'], batch['mask_padding'])
+
+        logits_observations = rearrange(outputs.logits_observations[:, :-1], 'b t o -> (b t) o')
+        loss_obs = F.cross_entropy(logits_observations, labels_observations)
+        loss_rewards = F.cross_entropy(rearrange(outputs.logits_rewards, 'b t e -> (b t) e'), labels_rewards)
+        loss_ends = F.cross_entropy(rearrange(outputs.logits_ends, 'b t e -> (b t) e'), labels_ends)
+
+        return LossWithIntermediateLosses(loss_obs=loss_obs, loss_rewards=loss_rewards, loss_ends=loss_ends)
+
+    def compute_labels_world_model(self, obs_tokens: torch.Tensor, rewards: torch.Tensor, ends: torch.Tensor, mask_padding: torch.BoolTensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        assert torch.all(ends.sum(dim=1) <= 1)  # at most 1 done
+        mask_fill = torch.logical_not(mask_padding)
+        labels_observations = rearrange(obs_tokens.masked_fill(mask_fill.unsqueeze(-1).expand_as(obs_tokens), -100), 'b t k -> b (t k)')[:, 1:]
+        labels_rewards = (rewards.sign() + 1).masked_fill(mask_fill, -100).long()  # Rewards clipped to {-1, 0, 1}
+        labels_ends = ends.masked_fill(mask_fill, -100)
+        return labels_observations.reshape(-1), labels_rewards.reshape(-1), labels_ends.reshape(-1)
@@ -0,0 +1,44 @@
+from pathlib import Path
+
+import hydra
+from hydra.utils import instantiate
+from omegaconf import DictConfig
+import torch
+
+from agent import Agent
+from envs import SingleProcessEnv, WorldModelEnv
+from game import AgentEnv, EpisodeReplayEnv, Game
+from models.actor_critic import ActorCritic
+from models.world_model import WorldModel
+
+
+@hydra.main(config_path="../config", config_name="trainer")
+def main(cfg: DictConfig):
+    device = torch.device(cfg.common.device)
+    assert cfg.mode in ('world_model', 'episode_replay', 'agent')
+
+    if cfg.mode in ['world_model', 'agent']:
+        env_fn = lambda: instantiate(cfg.env.test)
+        test_env = SingleProcessEnv(env_fn)
+        tokenizer = instantiate(cfg.tokenizer)
+        world_model = WorldModel(obs_vocab_size=tokenizer.vocab_size, act_vocab_size=test_env.num_actions, config=instantiate(cfg.world_model))
+        if cfg.mode == 'world_model':
+            env = WorldModelEnv(tokenizer=tokenizer, world_model=world_model, pretrained_agent_path=Path('checkpoints/last.pt'), device=device, env=env_fn())
+            keymap = cfg.env.keymap
+        else:
+            actor_critic = ActorCritic(**cfg.actor_critic, act_vocab_size=test_env.num_actions)
+            agent = Agent(tokenizer, world_model, actor_critic).to(device)
+            agent.load(Path('checkpoints/last.pt'), device)
+            env = AgentEnv(agent, test_env, cfg.env.keymap)
+            keymap = 'empty'
+
+    else:
+        env = EpisodeReplayEnv(replay_keymap_name=cfg.env.keymap, episode_dir=Path('media/episodes'))
+        keymap = 'episode_replay'
+
+    game = Game(env, keymap_name=keymap, size=(600, 1200 if cfg.mode == 'agent' else 600), fps=cfg.fps, verbose=bool(cfg.header))
+    game.run()
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,283 @@
+from collections import defaultdict
+from pathlib import Path
+import shutil
+import sys
+import time
+from typing import Any, Dict, Optional, Tuple
+
+import hydra
+from hydra.utils import instantiate
+from omegaconf import DictConfig, OmegaConf
+import torch
+import torch.nn as nn
+from tqdm import tqdm
+import wandb
+
+from agent import Agent
+from collector import Collector
+from envs import SingleProcessEnv, MultiProcessEnv
+from episode import Episode
+from make_reconstructions import make_reconstructions_from_batch
+from models.actor_critic import ActorCritic
+from models.world_model import WorldModel
+from utils import configure_optimizer, EpisodeDirManager, set_seed
+
+
+class Trainer:
+    def __init__(self, cfg: DictConfig) -> None:
+        wandb.init(
+            config=OmegaConf.to_container(cfg, resolve=True),
+            reinit=True,
+            resume=True,
+            **cfg.wandb
+        )
+
+        if cfg.common.seed is not None:
+            set_seed(cfg.common.seed)
+
+        self.cfg = cfg
+        self.start_epoch = 1
+        self.device = torch.device(cfg.common.device)
+
+        self.ckpt_dir = Path('checkpoints')
+        self.media_dir = Path('media')
+        self.episode_dir = self.media_dir / 'episodes'
+        self.reconstructions_dir = self.media_dir / 'reconstructions'
+
+        if not cfg.common.resume:
+            config_dir = Path('config')
+            config_path = config_dir / 'trainer.yaml'
+            config_dir.mkdir(exist_ok=False, parents=False)
+            shutil.copy('.hydra/config.yaml', config_path)
+            wandb.save(str(config_path))
+            shutil.copytree(src=(Path(hydra.utils.get_original_cwd()) / "src"), dst="./src")
+            shutil.copytree(src=(Path(hydra.utils.get_original_cwd()) / "scripts"), dst="./scripts")
+            self.ckpt_dir.mkdir(exist_ok=False, parents=False)
+            self.media_dir.mkdir(exist_ok=False, parents=False)
+            self.episode_dir.mkdir(exist_ok=False, parents=False)
+            self.reconstructions_dir.mkdir(exist_ok=False, parents=False)
+
+        episode_manager_train = EpisodeDirManager(self.episode_dir / 'train', max_num_episodes=cfg.collection.train.num_episodes_to_save)
+        episode_manager_test = EpisodeDirManager(self.episode_dir / 'test', max_num_episodes=cfg.collection.test.num_episodes_to_save)
+        self.episode_manager_imagination = EpisodeDirManager(self.episode_dir / 'imagination', max_num_episodes=cfg.evaluation.actor_critic.num_episodes_to_save)
+
+        def create_env(cfg_env, num_envs):
+            env_fn = lambda: instantiate(cfg_env)
+            return MultiProcessEnv(env_fn, num_envs, should_wait_num_envs_ratio=1.0) if num_envs > 1 else SingleProcessEnv(env_fn)
+
+        if self.cfg.training.should:
+            train_env = create_env(cfg.env.train, cfg.collection.train.num_envs)
+            self.train_dataset = instantiate(cfg.datasets.train)
+            self.train_collector = Collector(train_env, self.train_dataset, episode_manager_train)
+
+        if self.cfg.evaluation.should:
+            test_env = create_env(cfg.env.test, cfg.collection.test.num_envs)
+            self.test_dataset = instantiate(cfg.datasets.test)
+            self.test_collector = Collector(test_env, self.test_dataset, episode_manager_test)
+
+        assert self.cfg.training.should or self.cfg.evaluation.should
+        env = train_env if self.cfg.training.should else test_env
+
+        tokenizer = instantiate(cfg.tokenizer)
+        world_model = WorldModel(obs_vocab_size=tokenizer.vocab_size, act_vocab_size=env.num_actions, config=instantiate(cfg.world_model))
+        actor_critic = ActorCritic(**cfg.actor_critic, act_vocab_size=env.num_actions)
+        self.agent = Agent(tokenizer, world_model, actor_critic).to(self.device)
+        print(f'{sum(p.numel() for p in self.agent.tokenizer.parameters())} parameters in agent.tokenizer')
+        print(f'{sum(p.numel() for p in self.agent.world_model.parameters())} parameters in agent.world_model')
+        print(f'{sum(p.numel() for p in self.agent.actor_critic.parameters())} parameters in agent.actor_critic')
+
+        self.optimizer_tokenizer = torch.optim.Adam(self.agent.tokenizer.parameters(), lr=cfg.training.learning_rate)
+        self.optimizer_world_model = configure_optimizer(self.agent.world_model, cfg.training.learning_rate, cfg.training.world_model.weight_decay)
+        self.optimizer_actor_critic = torch.optim.Adam(self.agent.actor_critic.parameters(), lr=cfg.training.learning_rate)
+
+        if cfg.initialization.path_to_checkpoint is not None:
+            self.agent.load(**cfg.initialization, device=self.device)
+
+        if cfg.common.resume:
+            self.load_checkpoint()
+
+    def run(self) -> None:
+
+        for epoch in range(self.start_epoch, 1 + self.cfg.common.epochs):
+
+            print(f"\nEpoch {epoch} / {self.cfg.common.epochs}\n")
+            start_time = time.time()
+            to_log = []
+
+            if self.cfg.training.should:
+                if epoch <= self.cfg.collection.train.stop_after_epochs:
+                    to_log += self.train_collector.collect(self.agent, epoch, **self.cfg.collection.train.config)
+                to_log += self.train_agent(epoch)
+
+            if self.cfg.evaluation.should and (epoch % self.cfg.evaluation.every == 0):
+                self.test_dataset.clear()
+                to_log += self.test_collector.collect(self.agent, epoch, **self.cfg.collection.test.config)
+                to_log += self.eval_agent(epoch)
+
+            if self.cfg.training.should:
+                self.save_checkpoint(epoch, save_agent_only=not self.cfg.common.do_checkpoint)
+
+            to_log.append({'duration': (time.time() - start_time) / 3600})
+            for metrics in to_log:
+                wandb.log({'epoch': epoch, **metrics})
+
+        self.finish()
+
+    def train_agent(self, epoch: int) -> None:
+        self.agent.train()
+        self.agent.zero_grad()
+
+        metrics_tokenizer, metrics_world_model, metrics_actor_critic = {}, {}, {}
+
+        cfg_tokenizer = self.cfg.training.tokenizer
+        cfg_world_model = self.cfg.training.world_model
+        cfg_actor_critic = self.cfg.training.actor_critic
+
+        w = self.cfg.training.sampling_weights
+
+        if epoch > cfg_tokenizer.start_after_epochs:
+            metrics_tokenizer = self.train_component(self.agent.tokenizer, self.optimizer_tokenizer, sequence_length=1, sample_from_start=True, sampling_weights=w, **cfg_tokenizer)
+        self.agent.tokenizer.eval()
+
+        if epoch > cfg_world_model.start_after_epochs:
+            metrics_world_model = self.train_component(self.agent.world_model, self.optimizer_world_model, sequence_length=self.cfg.common.sequence_length, sample_from_start=True, sampling_weights=w, tokenizer=self.agent.tokenizer, **cfg_world_model)
+        self.agent.world_model.eval()
+
+        if epoch > cfg_actor_critic.start_after_epochs:
+            metrics_actor_critic = self.train_component(self.agent.actor_critic, self.optimizer_actor_critic, sequence_length=1 + self.cfg.training.actor_critic.burn_in, sample_from_start=False, sampling_weights=w, tokenizer=self.agent.tokenizer, world_model=self.agent.world_model, **cfg_actor_critic)
+        self.agent.actor_critic.eval()
+
+        return [{'epoch': epoch, **metrics_tokenizer, **metrics_world_model, **metrics_actor_critic}]
+
+    def train_component(self, component: nn.Module, optimizer: torch.optim.Optimizer, steps_per_epoch: int, batch_num_samples: int, grad_acc_steps: int, max_grad_norm: Optional[float], sequence_length: int, sampling_weights: Optional[Tuple[float]], sample_from_start: bool, **kwargs_loss: Any) -> Dict[str, float]:
+        loss_total_epoch = 0.0
+        intermediate_losses = defaultdict(float)
+
+        for _ in tqdm(range(steps_per_epoch), desc=f"Training {str(component)}", file=sys.stdout):
+            optimizer.zero_grad()
+            for _ in range(grad_acc_steps):
+                batch = self.train_dataset.sample_batch(batch_num_samples, sequence_length, sampling_weights, sample_from_start)
+                batch = self._to_device(batch)
+
+                losses = component.compute_loss(batch, **kwargs_loss) / grad_acc_steps
+                loss_total_step = losses.loss_total
+                loss_total_step.backward()
+                loss_total_epoch += loss_total_step.item() / steps_per_epoch
+
+                for loss_name, loss_value in losses.intermediate_losses.items():
+                    intermediate_losses[f"{str(component)}/train/{loss_name}"] += loss_value / steps_per_epoch
+
+            if max_grad_norm is not None:
+                torch.nn.utils.clip_grad_norm_(component.parameters(), max_grad_norm)
+
+            optimizer.step()
+
+        metrics = {f'{str(component)}/train/total_loss': loss_total_epoch, **intermediate_losses}
+        return metrics
+
+    @torch.no_grad()
+    def eval_agent(self, epoch: int) -> None:
+        self.agent.eval()
+
+        metrics_tokenizer, metrics_world_model = {}, {}
+
+        cfg_tokenizer = self.cfg.evaluation.tokenizer
+        cfg_world_model = self.cfg.evaluation.world_model
+        cfg_actor_critic = self.cfg.evaluation.actor_critic
+
+        if epoch > cfg_tokenizer.start_after_epochs:
+            metrics_tokenizer = self.eval_component(self.agent.tokenizer, cfg_tokenizer.batch_num_samples, sequence_length=1)
+
+        if epoch > cfg_world_model.start_after_epochs:
+            metrics_world_model = self.eval_component(self.agent.world_model, cfg_world_model.batch_num_samples, sequence_length=self.cfg.common.sequence_length, tokenizer=self.agent.tokenizer)
+
+        if epoch > cfg_actor_critic.start_after_epochs:
+            self.inspect_imagination(epoch)
+
+        if cfg_tokenizer.save_reconstructions:
+            batch = self._to_device(self.test_dataset.sample_batch(batch_num_samples=3, sequence_length=self.cfg.common.sequence_length))
+            make_reconstructions_from_batch(batch, save_dir=self.reconstructions_dir, epoch=epoch, tokenizer=self.agent.tokenizer)
+
+        return [metrics_tokenizer, metrics_world_model]
+
+    @torch.no_grad()
+    def eval_component(self, component: nn.Module, batch_num_samples: int, sequence_length: int, **kwargs_loss: Any) -> Dict[str, float]:
+        loss_total_epoch = 0.0
+        intermediate_losses = defaultdict(float)
+
+        steps = 0
+        pbar = tqdm(desc=f"Evaluating {str(component)}", file=sys.stdout)
+        for batch in self.test_dataset.traverse(batch_num_samples, sequence_length):
+            batch = self._to_device(batch)
+
+            losses = component.compute_loss(batch, **kwargs_loss)
+            loss_total_epoch += losses.loss_total.item()
+
+            for loss_name, loss_value in losses.intermediate_losses.items():
+                intermediate_losses[f"{str(component)}/eval/{loss_name}"] += loss_value
+
+            steps += 1
+            pbar.update(1)
+
+        intermediate_losses = {k: v / steps for k, v in intermediate_losses.items()}
+        metrics = {f'{str(component)}/eval/total_loss': loss_total_epoch / steps, **intermediate_losses}
+        return metrics
+
+    @torch.no_grad()
+    def inspect_imagination(self, epoch: int) -> None:
+        mode_str = 'imagination'
+        batch = self.test_dataset.sample_batch(batch_num_samples=self.episode_manager_imagination.max_num_episodes, sequence_length=1 + self.cfg.training.actor_critic.burn_in, sample_from_start=False)
+        outputs = self.agent.actor_critic.imagine(self._to_device(batch), self.agent.tokenizer, self.agent.world_model, horizon=self.cfg.evaluation.actor_critic.horizon, show_pbar=True)
+
+        to_log = []
+        for i, (o, a, r, d) in enumerate(zip(outputs.observations.cpu(), outputs.actions.cpu(), outputs.rewards.cpu(), outputs.ends.long().cpu())):  # Make everything (N, T, ...) instead of (T, N, ...)
+            episode = Episode(o, a, r, d, torch.ones_like(d))
+            episode_id = (epoch - 1 - self.cfg.training.actor_critic.start_after_epochs) * outputs.observations.size(0) + i
+            self.episode_manager_imagination.save(episode, episode_id, epoch)
+
+            metrics_episode = {k: v for k, v in episode.compute_metrics().__dict__.items()}
+            metrics_episode['episode_num'] = episode_id
+            metrics_episode['action_histogram'] = wandb.Histogram(episode.actions.numpy(), num_bins=self.agent.world_model.act_vocab_size)
+            to_log.append({f'{mode_str}/{k}': v for k, v in metrics_episode.items()})
+
+        return to_log
+
+    def _save_checkpoint(self, epoch: int, save_agent_only: bool) -> None:
+        torch.save(self.agent.state_dict(), self.ckpt_dir / 'last.pt')
+        if not save_agent_only:
+            torch.save(epoch, self.ckpt_dir / 'epoch.pt')
+            torch.save({
+                "optimizer_tokenizer": self.optimizer_tokenizer.state_dict(),
+                "optimizer_world_model": self.optimizer_world_model.state_dict(),
+                "optimizer_actor_critic": self.optimizer_actor_critic.state_dict(),
+            }, self.ckpt_dir / 'optimizer.pt')
+            ckpt_dataset_dir = self.ckpt_dir / 'dataset'
+            ckpt_dataset_dir.mkdir(exist_ok=True, parents=False)
+            self.train_dataset.update_disk_checkpoint(ckpt_dataset_dir)
+            if self.cfg.evaluation.should:
+                torch.save(self.test_dataset.num_seen_episodes, self.ckpt_dir / 'num_seen_episodes_test_dataset.pt')
+
+    def save_checkpoint(self, epoch: int, save_agent_only: bool) -> None:
+        tmp_checkpoint_dir = Path('checkpoints_tmp')
+        shutil.copytree(src=self.ckpt_dir, dst=tmp_checkpoint_dir, ignore=shutil.ignore_patterns('dataset'))
+        self._save_checkpoint(epoch, save_agent_only)
+        shutil.rmtree(tmp_checkpoint_dir)
+
+    def load_checkpoint(self) -> None:
+        assert self.ckpt_dir.is_dir()
+        self.start_epoch = torch.load(self.ckpt_dir / 'epoch.pt') + 1
+        self.agent.load(self.ckpt_dir / 'last.pt', device=self.device)
+        ckpt_opt = torch.load(self.ckpt_dir / 'optimizer.pt', map_location=self.device)
+        self.optimizer_tokenizer.load_state_dict(ckpt_opt['optimizer_tokenizer'])
+        self.optimizer_world_model.load_state_dict(ckpt_opt['optimizer_world_model'])
+        self.optimizer_actor_critic.load_state_dict(ckpt_opt['optimizer_actor_critic'])
+        self.train_dataset.load_disk_checkpoint(self.ckpt_dir / 'dataset')
+        if self.cfg.evaluation.should:
+            self.test_dataset.num_seen_episodes = torch.load(self.ckpt_dir / 'num_seen_episodes_test_dataset.pt')
+        print(f'Successfully loaded model, optimizer and {len(self.train_dataset)} episodes from {self.ckpt_dir.absolute()}.')
+
+    def _to_device(self, batch: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        return {k: batch[k].to(self.device) for k in batch}
+
+    def finish(self) -> None:
+        wandb.finish()
@@ -0,0 +1,142 @@
+from collections import OrderedDict
+from pathlib import Path
+import random
+import shutil
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from episode import Episode
+
+
+def configure_optimizer(model, learning_rate, weight_decay, *blacklist_module_names):
+    """Credits to https://github.com/karpathy/minGPT"""
+    # separate out all parameters to those that will and won't experience regularizing weight decay
+    decay = set()
+    no_decay = set()
+    whitelist_weight_modules = (torch.nn.Linear, torch.nn.Conv1d)
+    blacklist_weight_modules = (torch.nn.LayerNorm, torch.nn.Embedding)
+    for mn, m in model.named_modules():
+        for pn, p in m.named_parameters():
+            fpn = '%s.%s' % (mn, pn) if mn else pn  # full param name
+            if any([fpn.startswith(module_name) for module_name in blacklist_module_names]):
+                no_decay.add(fpn)
+            elif 'bias' in pn:
+                # all biases will not be decayed
+                no_decay.add(fpn)
+            elif pn.endswith('weight') and isinstance(m, whitelist_weight_modules):
+                # weights of whitelist modules will be weight decayed
+                decay.add(fpn)
+            elif pn.endswith('weight') and isinstance(m, blacklist_weight_modules):
+                # weights of blacklist modules will NOT be weight decayed
+                no_decay.add(fpn)
+
+    # validate that we considered every parameter
+    param_dict = {pn: p for pn, p in model.named_parameters()}
+    inter_params = decay & no_decay
+    union_params = decay | no_decay
+    assert len(inter_params) == 0, f"parameters {str(inter_params)} made it into both decay/no_decay sets!"
+    assert len(param_dict.keys() - union_params) == 0, f"parameters {str(param_dict.keys() - union_params)} were not separated into either decay/no_decay set!"
+
+    # create the pytorch optimizer object
+    optim_groups = [
+        {"params": [param_dict[pn] for pn in sorted(list(decay))], "weight_decay": weight_decay},
+        {"params": [param_dict[pn] for pn in sorted(list(no_decay))], "weight_decay": 0.0},
+    ]
+    optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate)
+    return optimizer
+
+
+def init_weights(module):
+    if isinstance(module, (nn.Linear, nn.Embedding)):
+        module.weight.data.normal_(mean=0.0, std=0.02)
+        if isinstance(module, nn.Linear) and module.bias is not None:
+            module.bias.data.zero_()
+    elif isinstance(module, nn.LayerNorm):
+        module.bias.data.zero_()
+        module.weight.data.fill_(1.0)
+
+
+def extract_state_dict(state_dict, module_name):
+    return OrderedDict({k.split('.', 1)[1]: v for k, v in state_dict.items() if k.startswith(module_name)})
+
+
+def set_seed(seed):
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    random.seed(seed)
+
+
+def remove_dir(path, should_ask=False):
+    assert path.is_dir()
+    if (not should_ask) or input(f"Remove directory : {path} ? [Y/n] ").lower() != 'n':
+        shutil.rmtree(path)
+
+
+def compute_lambda_returns(rewards, values, ends, gamma, lambda_):
+    assert rewards.ndim == 2 or (rewards.ndim == 3 and rewards.size(2) == 1)
+    assert rewards.shape == ends.shape == values.shape, f"{rewards.shape}, {values.shape}, {ends.shape}"  # (B, T, 1)
+    t = rewards.size(1)
+    lambda_returns = torch.empty_like(values)
+    lambda_returns[:, -1] = values[:, -1]
+    lambda_returns[:, :-1] = rewards[:, :-1] + ends[:, :-1].logical_not() * gamma * (1 - lambda_) * values[:, 1:]
+
+    last = values[:, -1]
+    for i in list(range(t - 1))[::-1]:
+        lambda_returns[:, i] += ends[:, i].logical_not() * gamma * lambda_ * last
+        last = lambda_returns[:, i]
+
+    return lambda_returns
+
+
+class LossWithIntermediateLosses:
+    def __init__(self, **kwargs):
+        self.loss_total = sum(kwargs.values())
+        self.intermediate_losses = {k: v.item() for k, v in kwargs.items()}
+
+    def __truediv__(self, value):
+        for k, v in self.intermediate_losses.items():
+            self.intermediate_losses[k] = v / value
+        self.loss_total = self.loss_total / value
+        return self
+
+
+class EpisodeDirManager:
+    def __init__(self, episode_dir: Path, max_num_episodes: int) -> None:
+        self.episode_dir = episode_dir
+        self.episode_dir.mkdir(parents=False, exist_ok=True)
+        self.max_num_episodes = max_num_episodes
+        self.best_return = float('-inf')
+
+    def save(self, episode: Episode, episode_id: int, epoch: int) -> None:
+        if self.max_num_episodes is not None and self.max_num_episodes > 0:
+            self._save(episode, episode_id, epoch)
+
+    def _save(self, episode: Episode, episode_id: int, epoch: int) -> None:
+        ep_paths = [p for p in self.episode_dir.iterdir() if p.stem.startswith('episode_')]
+        assert len(ep_paths) <= self.max_num_episodes
+        if len(ep_paths) == self.max_num_episodes:
+            to_remove = min(ep_paths, key=lambda ep_path: int(ep_path.stem.split('_')[1]))
+            to_remove.unlink()
+        episode.save(self.episode_dir / f'episode_{episode_id}_epoch_{epoch}.pt')
+
+        ep_return = episode.compute_metrics().episode_return
+        if ep_return > self.best_return:
+            self.best_return = ep_return
+            path_best_ep = [p for p in self.episode_dir.iterdir() if p.stem.startswith('best_')]
+            assert len(path_best_ep) in (0, 1)
+            if len(path_best_ep) == 1:
+                path_best_ep[0].unlink()
+            episode.save(self.episode_dir / f'best_episode_{episode_id}_epoch_{epoch}.pt')
+
+
+class RandomHeuristic:
+    def __init__(self, num_actions):
+        self.num_actions = num_actions
+
+    def act(self, obs):
+        assert obs.ndim == 4  # (N, H, W, C)
+        n = obs.size(0)
+        return torch.randint(low=0, high=self.num_actions, size=(n,))
				`@@ -0,0 +1 @@`
				{"Alien": [808.5], "Amidar": [148.6], "Assault": [1263.1], "Asterix": [25557.8], "BankHeist": [351.0], "BattleZone": [13871.2], "Boxing": [52.7], "Breakout": [414.1], "ChopperCommand": [1117.3], "CrazyClimber": [83940.2], "DemonAttack": [13003.9], "Freeway": [21.8], "Frostbite": [296.3], "Gopher": [3260.3], "Hero": [9315.9], "Jamesbond": [517.0], "Kangaroo": [724.1], "Krull": [5663.3], "KungFuMaster": [30944.8], "MsPacman": [1281.2], "Pong": [20.1], "PrivateEye": [96.7], "Qbert": [13781.9], "RoadRunner": [17751.3], "Seaquest": [1100.2], "UpNDown": [17264.2]}
				`@@ -0,0 +1 @@`
				{"Alien": [7127.7], "Amidar": [1719.5], "Assault": [742.0], "Asterix": [8503.3], "BankHeist": [753.1], "BattleZone": [37187.5], "Boxing": [12.1], "Breakout": [30.5], "ChopperCommand": [7387.8], "CrazyClimber": [35829.4], "DemonAttack": [1971.0], "Freeway": [29.6], "Frostbite": [4334.7], "Gopher": [2412.5], "Hero": [30826.4], "Jamesbond": [302.8], "Kangaroo": [3035.0], "Krull": [2665.5], "KungFuMaster": [22736.3], "MsPacman": [6951.6], "Pong": [14.6], "PrivateEye": [69571.3], "Qbert": [13455.0], "RoadRunner": [7845.0], "Seaquest": [42054.7], "UpNDown": [11693.2]}
				`@@ -0,0 +1 @@`
				{"Alien": [530.0], "Amidar": [38.8], "Assault": [500.1], "Asterix": [1734.0], "BankHeist": [192.5], "BattleZone": [7687.5], "Boxing": [15.1], "Breakout": [48.0], "ChopperCommand": [1350.0], "CrazyClimber": [56937.0], "DemonAttack": [3527.0], "Freeway": [21.8], "Frostbite": [255.0], "Gopher": [1256.0], "Hero": [3095.0], "Jamesbond": [87.5], "Kangaroo": [62.5], "Krull": [4890.8], "KungFuMaster": [18813.0], "MsPacman": [1265.6], "Pong": [-6.7], "PrivateEye": [56.3], "Qbert": [3952.0], "RoadRunner": [2500.0], "Seaquest": [208.0], "UpNDown": [2896.9]}
				`@@ -0,0 +1 @@`
				{"Alien": [227.8], "Amidar": [5.8], "Assault": [222.4], "Asterix": [210.0], "BankHeist": [14.2], "BattleZone": [2360.0], "Boxing": [0.1], "Breakout": [1.7], "ChopperCommand": [811.0], "CrazyClimber": [10780.5], "DemonAttack": [152.1], "Freeway": [0.0], "Frostbite": [65.2], "Gopher": [257.6], "Hero": [1027.0], "Jamesbond": [29.0], "Kangaroo": [52.0], "Krull": [1598.0], "KungFuMaster": [258.5], "MsPacman": [307.3], "Pong": [-20.7], "PrivateEye": [24.9], "Qbert": [163.9], "RoadRunner": [11.5], "Seaquest": [68.4], "UpNDown": [533.4]}
				`@@ -0,0 +1 @@`
				{"Alien": [677.1875, 435.46875, 515.9375, 425.3125, 1030.46875], "Amidar": [93.65625, 102.703125, 63.71875, 31.0, 80.203125], "Assault": [457.890625, 606.1875, 509.0625, 671.078125, 391.78125], "Asterix": [1061.71875, 1108.59375, 920.3125, 1065.625, 1485.15625], "BankHeist": [43.90625, 77.5, 0.0, 31.875, 17.96875], "BattleZone": [4031.2, 4031.2, 5359.375, 3484.375, 3250.0], "Boxing": [12.65625, 20.984375, 0.09375, -4.359375, 9.375], "Breakout": [10.3125, 16.4, 22.765625, 16.4, 16.046875], "ChopperCommand": [962.5, 896.875, 909.375, 848.4375, 1279.6875], "CrazyClimber": [54370.3125, 57220.3125, 51150.0, 87593.75, 62583.6], "DemonAttack": [288.90625, 208.1, 161.25, 177.578125, 204.84375], "Freeway": [0.0, 0.015625, 21.546875, 29.890625, 31.96875], "Frostbite": [262.34375, 248.75, 205.78125, 266.875, 200.625], "Gopher": [447.5, 732.5, 525.3125, 432.8125, 845.625], "Hero": [3061.5625, 3020.859375, 2060.390625, 2935.0, 2205.234375], "Jamesbond": [141.40625, 50.0, 110.9375, 123.4375, 76.5625], "Kangaroo": [65.625, 62.5, 25.0, 40.625, 62.5], "Krull": [3317.546875, 1933.71875, 1739.03125, 1379.625, 2654.28125], "KungFuMaster": [14318.75, 15023.4375, 14134.375, 21054.6875, 9781.25], "MsPacman": [1551.09375, 1221.25, 1199.0625, 1525.0, 1903.4375], "Pong": [21.0, -17.890625, 19.3125, 20.671875, 20.984375], "PrivateEye": [100.0, 0.0, 100.0, -25.140625, 0.0], "Qbert": [4258.984375, 150.0, 580.078125, 798.828125, 656.25], "RoadRunner": [3081.25, 900.0, 5906.25, 11176.5625, 7139.0625], "Seaquest": [711.5625, 455.3125, 565.625, 854.0625, 830.0], "UpNDown": [8614.53125, 1616.09375, 498.59375, 661.25, 5360.9375]}
				`@@ -0,0 +1 @@`
				`python src/main.py common.resume=True hydra.output_subdir=null hydra.run.dir=.`
				`@@ -0,0 +1 @@`
				`from .transformer import Transformer, TransformerConfig`