Training a Simple World Model with Jax

You used the existing models or fine-tuned the models for your case? I tried to build one with mamba ssm on my own I am not really good with codes I tired ai assisted path couldnt verify what I was doing but this sounds interesting.