Views
No views yet
1from disco_torch import DiscoTrainer, collect_rollout
2
3agent = YourAgent(obs_dim=64, num_actions=3).to(device)
4trainer = DiscoTrainer(agent, device=device) # auto-downloads weights
5
6env = YourEnv(num_envs=2)
7obs = env.obs()
8lstm_state = agent.init_lstm_state(env.num_envs, device)
9
10def step_fn(actions):
11 rewards, dones = env.step(actions)
12 return env.obs(), rewards, dones
13
14for step in range(1000):
15 rollout, obs, lstm_state = collect_rollout(
16 agent, step_fn, obs, lstm_state, rollout_len=29, device=device,
17 )
18 logs = trainer.step(rollout) # replay buffer, gradient loop, target updates — all handled
19
20DiscoTrainer encapsulates the full training loop: replay buffer, 32x inner gradient steps, per-element gradient
21clipping, Polyak target network updates, and meta-state management. See
22https://github.com/asystemoffields/disco-torch/blob/main/examples/catch_disco.py for a complete working example that
23reaches 99% catch rate in 1000 steps.
24
25Advanced: Low-level API
26
27from disco_torch import DiscoUpdateRule, load_disco103_weights
28
29rule = DiscoUpdateRule()
30load_disco103_weights(rule) # auto-downloads from this repo
31
32# Generate loss targets from a rollout
33meta_out, new_state = rule.unroll_meta_net(
34 rollout, agent_params, meta_state, unroll_fn, hyper_params
35)
36loss, logs = rule.agent_loss(rollout, meta_out, hyper_params)
37
38File
39
40disco_103.npz — NumPy archive with 42 parameters (754,778 values total), converted from the original JAX checkpoint.
41
42PyTorch Port
43
44See https://github.com/asystemoffields/disco-torch for the full PyTorch implementation, examples, and experiment
45results.
46
47Citation
48
49@article{oh2025disco,
50 title={Discovering State-of-the-art Reinforcement Learning Algorithms},
51 author={Oh, Junhyuk and Farquhar, Greg and Kemaev, Iurii and Calian, Dan A. and Hessel, Matteo and Zintgraf, Luisa
52and Singh, Satinder and van Hasselt, Hado and Silver, David},
53 journal={Nature},
54 volume={648},
55 pages={312--319},
56 year={2025},
57 doi={10.1038/s41586-025-09761-x}
58}
59
60License
61
62Apache 2.0 — same as the original https://github.com/google-deepmind/disco_rl.