[Probe Memory/Speed Issues] - Add scikit-learn based probe

mila-iqia / atari-representation-learning

Code for "Unsupervised State Representation Learning in Atari"

https://arxiv.org/abs/1906.08226

MIT License

240 stars 50 forks source link

[Probe Memory/Speed Issues] - Add scikit-learn based probe #46

Closed eracah closed 5 years ago

eracah commented 5 years ago

I had some memory issues with the probe as is when using AtariARI for another project. Instead of training a linear classifier in pytorch, it's a bit neater, cleaner, faster, more memory efficient to use scikit-learn for this purpose.