contrastive-rl-pytorch

contrastive-rl

For following a new line of research that started in 2022 from Eysenbach et al.

This is important not because of contrastive learning, but because it happens to be a special case where the RL and SSL algorithm is one. It reveals how "traditional" RL is unable to build up representations alone.

Update: Finally seeing it, at about 3-5k steps

install

$ pip install contrastive-rl-pytorch

usage

import torch
from contrastive_rl_pytorch import ContrastiveRLTrainer

from x_mlps_pytorch import ResidualNormedMLP # https://arxiv.org/abs/2503.14858

encoder = ResidualNormedMLP(dim = 256, dim_in = 16, dim_out = 128, keel_post_ln = True)

trainer = ContrastiveRLTrainer(encoder)

trajectories = torch.randn(256, 512, 16)

trainer(trajectories, 100)

# train for 100 steps and save

torch.save(encoder.state_dict(), './trained.pt')

quick test

make sure uv is installed pip install uv

then

$ uv run train_lunar.py --cpu

wait until 3-5k steps at least

citations

@misc{eysenbach2023contrastivelearninggoalconditionedreinforcement,
    title   = {Contrastive Learning as Goal-Conditioned Reinforcement Learning},
    author  = {Benjamin Eysenbach and Tianjun Zhang and Ruslan Salakhutdinov and Sergey Levine},
    year    = {2023},
    eprint  = {2206.07568},
    archivePrefix = {arXiv},
    primaryClass = {cs.LG},
    url     = {https://arxiv.org/abs/2206.07568},
}

@misc{ziarko2025contrastiverepresentationstemporalreasoning,
    title   = {Contrastive Representations for Temporal Reasoning},
    author  = {Alicja Ziarko and Michal Bortkiewicz and Michal Zawalski and Benjamin Eysenbach and Piotr Milos},
    year    = {2025},
    eprint  = {2508.13113},
    archivePrefix = {arXiv},
    primaryClass = {cs.LG},
    url     = {https://arxiv.org/abs/2508.13113},
}

@inproceedings{anonymous2025hierarchical,
    title   = {Hierarchical Contrastive Reinforcement Learning: learn representation more suitable for {RL} environments},
    author  = {Anonymous},
    booktitle = {Submitted to The Fourteenth International Conference on Learning Representations},
    year    = {2025},
    url     = {https://openreview.net/forum?id=rTCSFOzVcK},
    note    = {under review}
}

@misc{liu2024singlegoalneedskills,
    title   = {A Single Goal is All You Need: Skills and Exploration Emerge from Contrastive RL without Rewards, Demonstrations, or Subgoals},
    author  = {Grace Liu and Michael Tang and Benjamin Eysenbach},
    year    = {2024},
    eprint  = {2408.05804},
    archivePrefix = {arXiv},
    primaryClass = {cs.LG},
    url     = {https://arxiv.org/abs/2408.05804},
}

@inproceedings{anonymous2025demystifying,
    title   = {Demystifying Emergent Exploration in Goal-Conditioned {RL}},
    author  = {Anonymous},
    booktitle = {Submitted to The Fourteenth International Conference on Learning Representations},
    year    = {2025},
    url     = {https://openreview.net/forum?id=mwgYORsqtv},
    note    = {under review}
}

@inproceedings{wang2025,
    title   = {1000 Layer Networks for Self-Supervised {RL}: Scaling Depth Can Enable New Goal-Reaching Capabilities},
    author  = {Kevin Wang and Ishaan Javali and Micha{\l} Bortkiewicz and Tomasz Trzcinski and Benjamin Eysenbach},
    booktitle = {The Thirty-ninth Annual Conference on Neural Information Processing Systems},
    year    = {2025},
    url     = {https://openreview.net/forum?id=s0JVsx3bx1}
}

@misc{nimonkar2025selfsupervisedgoalreachingresultsmultiagent,
    title   = {Self-Supervised Goal-Reaching Results in Multi-Agent Cooperation and Exploration},
    author  = {Chirayu Nimonkar and Shlok Shah and Catherine Ji and Benjamin Eysenbach},
    year    = {2025},
    eprint  = {2509.10656},
    archivePrefix = {arXiv},
    primaryClass = {cs.LG},
    url     = {https://arxiv.org/abs/2509.10656},
}

@misc{yin2026emergentdexteritydiverseresets,
    title   = {Emergent Dexterity via Diverse Resets and Large-Scale Reinforcement Learning},
    author  = {Patrick Yin and Tyler Westenbroek and Zhengyu Zhang and Joshua Tran and Ignacio Dagnino and Eeshani Shilamkar and Numfor Mbiziwo-Tiapo and Simran Bagaria and Xinlei Liu and Galen Mullins and Andrey Kolobov and Abhishek Gupta},
    year    = {2026},
    eprint  = {2603.15789},
    archivePrefix = {arXiv},
    primaryClass = {cs.RO},
    url     = {https://arxiv.org/abs/2603.15789},
}

@inproceedings{wang2023optimal,
    title   = {Optimal Goal-Reaching Reinforcement Learning via Quasimetric Learning},
    author  = {Tongzhou Wang and Antonio Torralba and Phillip Isola and Amy Zhang},
    booktitle = {International Conference on Machine Learning (ICML)},
    year    = {2023},
    url     = {https://arxiv.org/abs/2304.01203}
}

@misc{korniak2026stepstimelearningrepresentations,
    title   = {Three Steps at a Time: Learning Representations from Action Sequences in Contrastive RL},
    author  = {Michal Korniak and Kamil Dybek and Benjamin Eysenbach and Marco Bagatella and Micha{\l} Bortkiewicz},
    year    = {2026},
    eprint  = {2608.30640},
    archivePrefix = {arXiv},
    primaryClass = {cs.LG},
    url     = {https://arxiv.org/abs/2608.30640}
}

Release files for contrastive-rl-pytorch 0.4.7

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for contrastive-rl-pytorch 0.4.7
File	Size	Uploaded
contrastive_rl_pytorch-0.4.7.tar.gz	5.8 kB	Sep 10, 2026	Details

Built distribution (wheel)

Table of built distributions (wheels) for contrastive-rl-pytorch 0.4.7
File	Interpreter	ABI	Platform	Reset
contrastive_rl_pytorch-0.4.7-py3-none-any.whl 4.7 kB Sep 10, 2026	Python 3	none	any	Details

Total release size: 10.6 kB

Release files / contrastive_rl_pytorch-0.4.7.tar.gz

Download URL	contrastive_rl_pytorch-0.4.7.tar.gz
Size	5.8 kB
Tags	Source
SHA-256 checksum How to use checksums	`f48c0fadf77592a21df494104c0552e73b3056b442d95f1917a6cb72be0d8fbf`
BLAKE2b-256 checksum How to use checksums	`921654a7f5490f3e0fe95bbf529346c5638ff9997d58ef533ac483c47f765819`
Upload date	Sep 10, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`uv/0.8.17`

Release files / contrastive_rl_pytorch-0.4.7-py3-none-any.whl

Download URL	contrastive_rl_pytorch-0.4.7-py3-none-any.whl
Size	4.7 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`74cbbee7f061e31c4663025d14a96491d970f755820c84fe5d0d251cedeb9fa0`
BLAKE2b-256 checksum How to use checksums	`dd07ac87c6ef918d12ee59d5e555244d7b71124bb30c53e91ff737aa2e3d743a`
Upload date	Sep 10, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`uv/0.8.17`