EdanToledo
commented
9 months ago I have done this before and can confirm this is possible and can bring speed improvements. It might not be so easy to modify the dreamerv3 open-source codebase but doing it from scratch using the codebase as a starting point is worthwhile. I'm planning on doing it again sometime soon but not sure when I'll get around to it. I'm interested in reviewing or collaborating if you want to do it.
DreamerV3 is a SOTA MbRL algorithm. It's already written in jax. I want to gather if anyone here is potentially interested to make it's environment interactions also into jax. ie If you have a purejax environment, can we modify the current dreamerv3 code so that we have a blazing fast setup similar to other algorithms here. I am just warming up to jax, so I don't have too much insight into it yet.