disco-torch 0.1.0__tar.gz

disco_torch-0.1.0/LICENSE

@@ -0,0 +1,191 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to the Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by the Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding any notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   Copyright 2025 asystemoffields
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

disco_torch-0.1.0/PKG-INFO

@@ -0,0 +1,250 @@
+Metadata-Version: 2.4
+Name: disco-torch
+Version: 0.1.0
+Summary: PyTorch port of DeepMind's Disco103 meta-learned RL update rule
+Author: asystemoffields
+License-Expression: Apache-2.0
+Project-URL: Homepage, https://github.com/asystemoffields/disco-torch
+Project-URL: Repository, https://github.com/asystemoffields/disco-torch
+Project-URL: Issues, https://github.com/asystemoffields/disco-torch/issues
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.0
+Requires-Dist: numpy>=1.24
+Provides-Extra: hub
+Requires-Dist: huggingface-hub>=0.20; extra == "hub"
+Provides-Extra: examples
+Requires-Dist: gymnasium>=1.0; extra == "examples"
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: gymnasium>=1.0; extra == "dev"
+Requires-Dist: huggingface-hub>=0.20; extra == "dev"
+Provides-Extra: jax
+Requires-Dist: jax; extra == "jax"
+Requires-Dist: dm-haiku; extra == "jax"
+Requires-Dist: rlax; extra == "jax"
+Requires-Dist: distrax; extra == "jax"
+Dynamic: license-file
+
+# disco-torch
+
+A PyTorch port of DeepMind's **Disco103** — the meta-learned reinforcement learning update rule from [*Discovering State-of-the-art Reinforcement Learning Algorithms*](https://doi.org/10.1038/s41586-025-09761-x) (Nature, 2025).
+
+## What is DiscoRL?
+
+Instead of hand-crafted loss functions like PPO or GRPO, DiscoRL uses a small LSTM neural network (the "meta-network") that **generates loss targets** for RL agents. Given a rollout of agent experience — policy logits, rewards, advantages, auxiliary predictions — the meta-network outputs target distributions. The agent then minimizes KL divergence between its outputs and these learned targets.
+
+The Disco103 checkpoint (754,778 parameters) was meta-trained by DeepMind across thousands of Atari-like environments. It generalizes as a drop-in update rule for new tasks — no reward shaping, no hyperparameter-specific loss design.
+
+## Why a PyTorch port?
+
+The [original implementation](https://github.com/google-deepmind/disco_rl) uses JAX + Haiku. This port enables using Disco103 in PyTorch training pipelines without any JAX dependency at inference time.
+
+## Installation
+
+```bash
+pip install disco-torch
+```
+
+With optional extras:
+
+```bash
+pip install disco-torch[hub]       # HuggingFace Hub weight downloads
+pip install disco-torch[examples]  # gymnasium for running examples
+pip install disco-torch[dev]       # pytest + all extras for development
+```
+
+### Weights
+
+Option 1 — Download from HuggingFace Hub (requires `pip install disco-torch[hub]`):
+
+```python
+from disco_torch import load_disco103_weights
+
+rule = DiscoUpdateRule()
+load_disco103_weights(rule)  # auto-downloads from HuggingFace Hub
+```
+
+Option 2 — Manual download from the [disco_rl repo](https://github.com/google-deepmind/disco_rl):
+
+```bash
+cp path/to/disco_103.npz weights/
+```
+
+```python
+load_disco103_weights(rule, "weights/disco_103.npz")
+```
+
+## Quick start
+
+```python
+import torch
+from disco_torch import DiscoUpdateRule, UpdateRuleInputs, load_disco103_weights
+
+# Load the meta-network with pretrained weights
+rule = DiscoUpdateRule()
+load_disco103_weights(rule, "weights/disco_103.npz")
+
+# Initialize meta-RNN state (persists across training steps)
+state = rule.meta_net.initial_meta_rnn_state()
+
+# Run the meta-network on a rollout
+with torch.no_grad():
+    meta_out, new_state = rule.meta_net(inputs, state)
+    # meta_out["pi"]  — policy loss targets  [T, B, A]
+    # meta_out["y"]   — value loss targets   [T, B, 600]
+    # meta_out["z"]   — auxiliary loss targets [T, B, 600]
+```
+
+### Full training loop
+
+```python
+# At each learner step:
+meta_out, new_meta_state = rule.unroll_meta_net(
+    rollout, agent_params, meta_state, unroll_fn, hyper_params
+)
+
+# Compute agent loss (KL divergence against meta-network targets)
+loss, logs = rule.agent_loss(rollout, meta_out, hyper_params)
+
+# Value function loss (no meta-gradient)
+value_loss, value_logs = rule.agent_loss_no_meta(rollout, meta_out, hyper_params)
+```
+
+## Architecture
+
+```
+Outer (per-trajectory):
+  y_net           MLP [600 -> 16 -> 1]         Value prediction embedding
+  z_net           MLP [600 -> 16 -> 1]         Auxiliary prediction embedding
+  policy_net      Conv1dNet [9 -> 16 -> 2]     Action-conditional embedding
+  trajectory_rnn  LSTM(27, 256)                Reverse-unrolled over trajectory
+  state_gate      Linear(128 -> 256)           Multiplicative gate from meta-RNN
+  y_head / z_head Linear(256 -> 600)           Loss targets for y and z
+  pi_conv + head  Conv1dNet [258 -> 16] -> 1   Policy loss target (per action)
+
+Meta-RNN (per-lifetime):
+  Separate y/z/policy nets, input MLP(29 -> 16), LSTMCell(16, 128)
+```
+
+The outer network processes each trajectory with a reverse-unrolled LSTM. The meta-RNN operates at a slower timescale — it sees batch-time averages and modulates the outer network via a multiplicative gate. This two-level architecture lets the update rule adapt its behavior over an agent's lifetime.
+
+## End-to-end example: Tiny Transformer
+
+The flagship example trains a **1-layer causal transformer** (~141K params) to generate strictly increasing digit sequences, using Disco103 as the RL update rule. No supervised data — the agent learns purely from a per-token reward signal.
+
+**[Open in Google Colab](examples/tinylm_disco_colab.ipynb)** (runs on free T4 GPU)
+
+Using curriculum learning (sequence lengths 4→6→8), the agent achieves **86% strictly increasing** on 8-token sequences, discovering sequences like `[0, 1, 5, 6, 5, 7, 8, 9]`. See [`experiments.md`](experiments.md) for the full research log.
+
+```bash
+# Standalone script (local GPU or CPU)
+python examples/tinylm_disco.py --weights weights/disco_103.npz
+```
+
+A simpler CartPole example is also provided for API reference:
+
+```bash
+python examples/cartpole_disco.py --weights weights/disco_103.npz
+```
+
+> **Note:** Disco103 was meta-trained on 103 complex environments (Atari, ProcGen, DMLab-30). Both CartPole and token generation are outside this distribution. See [`experiments.md`](experiments.md) for analysis of how this affects transfer.
+
+## Package structure
+
+```
+disco_torch/
+  __init__.py          Public API
+  types.py             Dataclasses: UpdateRuleInputs, MetaNetInputOption, ValueOuts, etc.
+  transforms.py        Input transforms and construct_input()
+  meta_net.py          DiscoMetaNet — the full LSTM meta-network
+  update_rule.py       DiscoUpdateRule — meta-net + value computation + loss
+  value_utils.py       V-trace, TD-error, advantage estimation, Q-values
+  utils.py             batch_lookup, signed_logp1, 2-hot encoding, EMA
+  load_weights.py      Maps JAX/Haiku NPZ keys -> PyTorch modules
+
+examples/
+  tinylm_disco.py          Train a tiny transformer with Disco103 (standalone)
+  tinylm_disco_colab.ipynb Google Colab notebook (recommended)
+  cartpole_disco.py        CartPole API reference example
+
+scripts/
+  inspect_disco103.py      Print NPZ weight names and shapes
+  validate_against_jax.py  Numerical comparison: PyTorch vs JAX reference
+
+tests/
+  test_utils.py            Unit tests for utility functions
+  test_building_blocks.py  Unit tests for network building blocks
+  test_meta_net.py         Snapshot tests for meta-network forward pass
+```
+
+## Numerical validation
+
+All outputs match the JAX reference implementation within float32 precision:
+
+| Output | Max diff | Status |
+|--------|----------|--------|
+| pi (policy targets) | < 1.3e-06 | PASS |
+| y (value targets) | < 1.3e-06 | PASS |
+| z (auxiliary targets) | < 1.3e-06 | PASS |
+| meta_input_emb | < 1.3e-06 | PASS |
+| meta_rnn_h | < 1.3e-06 | PASS |
+
+To run the test suite (no JAX required):
+
+```bash
+pip install disco-torch[dev]
+pytest
+```
+
+To run JAX cross-validation (requires JAX + disco_rl):
+
+```bash
+pip install disco_rl jax dm-haiku rlax distrax
+python scripts/validate_against_jax.py
+```
+
+## Key implementation details
+
+- **HaikuLSTMCell**: Haiku uses gate order `[i, g, f, o]` with a +1 forget gate bias, vs PyTorch's `[i, f, g, o]`. This is handled by a custom LSTM cell.
+- **Weight mapping**: The 42 JAX/Haiku parameters have nested path names (e.g., `lstm/~/meta_lstm/~unroll/mlp_2/~/linear_0/w`). `load_weights.py` maps every one to the correct PyTorch module.
+- **Conv1dBlock**: Each block concatenates per-action features with their mean across actions before the convolution — matching the JAX implementation's broadcast pattern.
+- **Value utilities**: V-trace, Retrace-style Q-value estimation, signed hyperbolic transforms, and 2-hot categorical encoding are all ported.
+
+## Requirements
+
+- Python >= 3.11
+- PyTorch >= 2.0
+- NumPy >= 1.24
+
+## License
+
+Apache 2.0 — same as the original disco_rl.
+
+## Citation
+
+If you use this port, please cite the original paper:
+
+```bibtex
+@article{oh2025disco,
+  title={Discovering State-of-the-art Reinforcement Learning Algorithms},
+  author={Oh, Junhyuk and Farquhar, Greg and Kemaev, Iurii and Calian, Dan A. and Hessel, Matteo and Zintgraf, Luisa and Singh, Satinder and van Hasselt, Hado and Silver, David},
+  journal={Nature},
+  volume={648},
+  pages={312--319},
+  year={2025},
+  doi={10.1038/s41586-025-09761-x}
+}
+```
+
+## Acknowledgments
+
+This is a community port of [google-deepmind/disco_rl](https://github.com/google-deepmind/disco_rl). All credit for the algorithm, architecture, and pretrained weights goes to the original authors.

disco_torch-0.1.0/README.md

@@ -0,0 +1,214 @@
+# disco-torch
+
+A PyTorch port of DeepMind's **Disco103** — the meta-learned reinforcement learning update rule from [*Discovering State-of-the-art Reinforcement Learning Algorithms*](https://doi.org/10.1038/s41586-025-09761-x) (Nature, 2025).
+
+## What is DiscoRL?
+
+Instead of hand-crafted loss functions like PPO or GRPO, DiscoRL uses a small LSTM neural network (the "meta-network") that **generates loss targets** for RL agents. Given a rollout of agent experience — policy logits, rewards, advantages, auxiliary predictions — the meta-network outputs target distributions. The agent then minimizes KL divergence between its outputs and these learned targets.
+
+The Disco103 checkpoint (754,778 parameters) was meta-trained by DeepMind across thousands of Atari-like environments. It generalizes as a drop-in update rule for new tasks — no reward shaping, no hyperparameter-specific loss design.
+
+## Why a PyTorch port?
+
+The [original implementation](https://github.com/google-deepmind/disco_rl) uses JAX + Haiku. This port enables using Disco103 in PyTorch training pipelines without any JAX dependency at inference time.
+
+## Installation
+
+```bash
+pip install disco-torch
+```
+
+With optional extras:
+
+```bash
+pip install disco-torch[hub]       # HuggingFace Hub weight downloads
+pip install disco-torch[examples]  # gymnasium for running examples
+pip install disco-torch[dev]       # pytest + all extras for development
+```
+
+### Weights
+
+Option 1 — Download from HuggingFace Hub (requires `pip install disco-torch[hub]`):
+
+```python
+from disco_torch import load_disco103_weights
+
+rule = DiscoUpdateRule()
+load_disco103_weights(rule)  # auto-downloads from HuggingFace Hub
+```
+
+Option 2 — Manual download from the [disco_rl repo](https://github.com/google-deepmind/disco_rl):
+
+```bash
+cp path/to/disco_103.npz weights/
+```
+
+```python
+load_disco103_weights(rule, "weights/disco_103.npz")
+```
+
+## Quick start
+
+```python
+import torch
+from disco_torch import DiscoUpdateRule, UpdateRuleInputs, load_disco103_weights
+
+# Load the meta-network with pretrained weights
+rule = DiscoUpdateRule()
+load_disco103_weights(rule, "weights/disco_103.npz")
+
+# Initialize meta-RNN state (persists across training steps)
+state = rule.meta_net.initial_meta_rnn_state()
+
+# Run the meta-network on a rollout
+with torch.no_grad():
+    meta_out, new_state = rule.meta_net(inputs, state)
+    # meta_out["pi"]  — policy loss targets  [T, B, A]
+    # meta_out["y"]   — value loss targets   [T, B, 600]
+    # meta_out["z"]   — auxiliary loss targets [T, B, 600]
+```
+
+### Full training loop
+
+```python
+# At each learner step:
+meta_out, new_meta_state = rule.unroll_meta_net(
+    rollout, agent_params, meta_state, unroll_fn, hyper_params
+)
+
+# Compute agent loss (KL divergence against meta-network targets)
+loss, logs = rule.agent_loss(rollout, meta_out, hyper_params)
+
+# Value function loss (no meta-gradient)
+value_loss, value_logs = rule.agent_loss_no_meta(rollout, meta_out, hyper_params)
+```
+
+## Architecture
+
+```
+Outer (per-trajectory):
+  y_net           MLP [600 -> 16 -> 1]         Value prediction embedding
+  z_net           MLP [600 -> 16 -> 1]         Auxiliary prediction embedding
+  policy_net      Conv1dNet [9 -> 16 -> 2]     Action-conditional embedding
+  trajectory_rnn  LSTM(27, 256)                Reverse-unrolled over trajectory
+  state_gate      Linear(128 -> 256)           Multiplicative gate from meta-RNN
+  y_head / z_head Linear(256 -> 600)           Loss targets for y and z
+  pi_conv + head  Conv1dNet [258 -> 16] -> 1   Policy loss target (per action)
+
+Meta-RNN (per-lifetime):
+  Separate y/z/policy nets, input MLP(29 -> 16), LSTMCell(16, 128)
+```
+
+The outer network processes each trajectory with a reverse-unrolled LSTM. The meta-RNN operates at a slower timescale — it sees batch-time averages and modulates the outer network via a multiplicative gate. This two-level architecture lets the update rule adapt its behavior over an agent's lifetime.
+
+## End-to-end example: Tiny Transformer
+
+The flagship example trains a **1-layer causal transformer** (~141K params) to generate strictly increasing digit sequences, using Disco103 as the RL update rule. No supervised data — the agent learns purely from a per-token reward signal.
+
+**[Open in Google Colab](examples/tinylm_disco_colab.ipynb)** (runs on free T4 GPU)
+
+Using curriculum learning (sequence lengths 4→6→8), the agent achieves **86% strictly increasing** on 8-token sequences, discovering sequences like `[0, 1, 5, 6, 5, 7, 8, 9]`. See [`experiments.md`](experiments.md) for the full research log.
+
+```bash
+# Standalone script (local GPU or CPU)
+python examples/tinylm_disco.py --weights weights/disco_103.npz
+```
+
+A simpler CartPole example is also provided for API reference:
+
+```bash
+python examples/cartpole_disco.py --weights weights/disco_103.npz
+```
+
+> **Note:** Disco103 was meta-trained on 103 complex environments (Atari, ProcGen, DMLab-30). Both CartPole and token generation are outside this distribution. See [`experiments.md`](experiments.md) for analysis of how this affects transfer.
+
+## Package structure
+
+```
+disco_torch/
+  __init__.py          Public API
+  types.py             Dataclasses: UpdateRuleInputs, MetaNetInputOption, ValueOuts, etc.
+  transforms.py        Input transforms and construct_input()
+  meta_net.py          DiscoMetaNet — the full LSTM meta-network
+  update_rule.py       DiscoUpdateRule — meta-net + value computation + loss
+  value_utils.py       V-trace, TD-error, advantage estimation, Q-values
+  utils.py             batch_lookup, signed_logp1, 2-hot encoding, EMA
+  load_weights.py      Maps JAX/Haiku NPZ keys -> PyTorch modules
+
+examples/
+  tinylm_disco.py          Train a tiny transformer with Disco103 (standalone)
+  tinylm_disco_colab.ipynb Google Colab notebook (recommended)
+  cartpole_disco.py        CartPole API reference example
+
+scripts/
+  inspect_disco103.py      Print NPZ weight names and shapes
+  validate_against_jax.py  Numerical comparison: PyTorch vs JAX reference
+
+tests/
+  test_utils.py            Unit tests for utility functions
+  test_building_blocks.py  Unit tests for network building blocks
+  test_meta_net.py         Snapshot tests for meta-network forward pass
+```
+
+## Numerical validation
+
+All outputs match the JAX reference implementation within float32 precision:
+
+| Output | Max diff | Status |
+|--------|----------|--------|
+| pi (policy targets) | < 1.3e-06 | PASS |
+| y (value targets) | < 1.3e-06 | PASS |
+| z (auxiliary targets) | < 1.3e-06 | PASS |
+| meta_input_emb | < 1.3e-06 | PASS |
+| meta_rnn_h | < 1.3e-06 | PASS |
+
+To run the test suite (no JAX required):
+
+```bash
+pip install disco-torch[dev]
+pytest
+```
+
+To run JAX cross-validation (requires JAX + disco_rl):
+
+```bash
+pip install disco_rl jax dm-haiku rlax distrax
+python scripts/validate_against_jax.py
+```
+
+## Key implementation details
+
+- **HaikuLSTMCell**: Haiku uses gate order `[i, g, f, o]` with a +1 forget gate bias, vs PyTorch's `[i, f, g, o]`. This is handled by a custom LSTM cell.
+- **Weight mapping**: The 42 JAX/Haiku parameters have nested path names (e.g., `lstm/~/meta_lstm/~unroll/mlp_2/~/linear_0/w`). `load_weights.py` maps every one to the correct PyTorch module.
+- **Conv1dBlock**: Each block concatenates per-action features with their mean across actions before the convolution — matching the JAX implementation's broadcast pattern.
+- **Value utilities**: V-trace, Retrace-style Q-value estimation, signed hyperbolic transforms, and 2-hot categorical encoding are all ported.
+
+## Requirements
+
+- Python >= 3.11
+- PyTorch >= 2.0
+- NumPy >= 1.24
+
+## License
+
+Apache 2.0 — same as the original disco_rl.
+
+## Citation
+
+If you use this port, please cite the original paper:
+
+```bibtex
+@article{oh2025disco,
+  title={Discovering State-of-the-art Reinforcement Learning Algorithms},
+  author={Oh, Junhyuk and Farquhar, Greg and Kemaev, Iurii and Calian, Dan A. and Hessel, Matteo and Zintgraf, Luisa and Singh, Satinder and van Hasselt, Hado and Silver, David},
+  journal={Nature},
+  volume={648},
+  pages={312--319},
+  year={2025},
+  doi={10.1038/s41586-025-09761-x}
+}
+```
+
+## Acknowledgments
+
+This is a community port of [google-deepmind/disco_rl](https://github.com/google-deepmind/disco_rl). All credit for the algorithm, architecture, and pretrained weights goes to the original authors.

disco_torch-0.1.0/disco_torch/__init__.py

@@ -0,0 +1,16 @@
+"""disco_torch: PyTorch port of DeepMind's Disco103 meta-learned RL update rule."""
+
+from disco_torch.types import UpdateRuleInputs, ValueOuts, EmaState, MetaNetInputOption, TransformConfig
+from disco_torch.update_rule import DiscoUpdateRule
+from disco_torch.load_weights import load_disco103_weights, download_disco103_weights
+
+__all__ = [
+    "DiscoUpdateRule",
+    "UpdateRuleInputs",
+    "ValueOuts",
+    "EmaState",
+    "MetaNetInputOption",
+    "TransformConfig",
+    "load_disco103_weights",
+    "download_disco103_weights",
+]