65 lines
2.2 KiB
Markdown
65 lines
2.2 KiB
Markdown
|
---
|
||
|
|
license: apache-2.0
|
||
|
|
base_model: Qwen/Qwen2.5-7B-Instruct
|
||
|
|
tags:
|
||
|
|
- world-model
|
||
|
|
- babyai
|
||
|
|
- reinforcement-learning
|
||
|
|
- model-based-rl
|
||
|
|
- sft
|
||
|
|
- lora
|
||
|
|
datasets:
|
||
|
|
- GGOSinon/babyai-world-model-sft
|
||
|
|
language:
|
||
|
|
- en
|
||
|
|
pipeline_tag: text-generation
|
||
|
|
---
|
||
|
|
|
||
|
|
# BabyAI World Model (Qwen2.5-7B SFT)
|
||
|
|
|
||
|
|
A world model for the BabyAI grid-world environment, fine-tuned from Qwen2.5-7B-Instruct using LoRA. This model predicts the next observation and available actions given the current state and the agent's action.
|
||
|
|
|
||
|
|
## Model Details
|
||
|
|
|
||
|
|
- **Base model**: Qwen2.5-7B-Instruct
|
||
|
|
- **Fine-tuning**: LoRA (40.4M trainable params, 0.53% of 7.66B total), merged after training
|
||
|
|
- **Training data**: [GGOSinon/babyai-world-model-sft](https://huggingface.co/datasets/GGOSinon/babyai-world-model-sft) (58K transitions, 1 epoch)
|
||
|
|
- **Training time**: ~5.5 hours on 1x A100 40GB
|
||
|
|
- **Final loss**: 0.023
|
||
|
|
|
||
|
|
## Performance (Done-Detection, 102 test cases)
|
||
|
|
|
||
|
|
| Model | Accuracy | Precision | Recall | FPR |
|
||
|
|
|---|---|---|---|---|
|
||
|
|
| Qwen2.5-7B zero-shot | 76.5% | 91.7% | 32.4% | 1.5% |
|
||
|
|
| **Qwen2.5-7B SFT (this model)** | **97.1%** | **100%** | **91.2%** | **0.0%** |
|
||
|
|
| Gemini 2.5 Flash zero-shot | 97.1% | 100% | 91.2% | 0.0% |
|
||
|
|
|
||
|
|
## Usage
|
||
|
|
|
||
|
|
```python
|
||
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||
|
|
|
||
|
|
model = AutoModelForCausalLM.from_pretrained("GGOSinon/babyai-world-model-7B-sft", torch_dtype="bfloat16").to("cuda")
|
||
|
|
tokenizer = AutoTokenizer.from_pretrained("GGOSinon/babyai-world-model-7B-sft")
|
||
|
|
|
||
|
|
messages = [
|
||
|
|
{"role": "system", "content": "You are a simulator for a grid-world environment called BabyAI..."},
|
||
|
|
{"role": "user", "content": "Goal: pick up the red box\n\nObservation:\n...\nAvailable actions: [...]\nAgent's action: pickup red box 1"}
|
||
|
|
]
|
||
|
|
|
||
|
|
inputs = tokenizer.apply_chat_template(messages, tokenize=True, return_tensors="pt", add_generation_prompt=True).to("cuda")
|
||
|
|
output = model.generate(inputs, max_new_tokens=300, do_sample=False)
|
||
|
|
print(tokenizer.decode(output[0][inputs.shape[1]:], skip_special_tokens=True))
|
||
|
|
```
|
||
|
|
|
||
|
|
## Output Format
|
||
|
|
|
||
|
|
The model outputs in this format:
|
||
|
|
```
|
||
|
|
<observation>next observation text</observation>
|
||
|
|
<available_actions>["action1", "action2", ...]</available_actions>
|
||
|
|
```
|
||
|
|
|
||
|
|
Task completion is indicated by "The task is completed." appended to the observation text.
|