Retrieval-Augmented World Model

The current environments considered are:

1. BlocksWorld
2. Game24
3. BabyAI

1. run datacollecting/main_collecting_bw.py to collect of BlocksWorld
2. As well as for game24, and babyai, please run the respective files

Given the pre-trained embedding model and the LLM for prediction, we plug a MLP into the embedding model and train the MLP for better embedding.

To improve the training efficiency, we will consider using the smaller LLM for training and test the embedding model on other LLMs to demonstrate the generalizability of the learned embedding.