README.md

# RLWF — DreamZero checkpoints

Private checkpoint repository for the RLWF paper ("Active Robot Data Collection
from World Model Feedback"). Two checkpoints, both **stock DreamZero
architecture, no architectural modifications** — only the training data and
training-config differ.

## Layout

```
rlwf-ckpt/
├── README.md
├── LICENSE
├── mimicgen-core-14b-lora-step80000/   # LoRA fine-tune, ~217 MB
└── mimicgen-core-14b-full-step46000/   # full fine-tune, 10-shard ~47 GB
```

## What each checkpoint is

### `mimicgen-core-14b-lora-step80000/`

- **Architecture**: stock DreamZero (`groot.vla.model.dreamzero.base_vla.VLA`)
- **Base model**: Wan2.1-I2V-14B-480P, frozen
- **Adapter**: LoRA, rank 4, target modules `q,k,v,o,ffn.0,ffn.2`
- **Action head**: WAN flow-matching action transformer
  (`groot.vla.model.dreamzero.action_head.wan_flow_matching_action_tf.WANPolicyHead`)
- **Action dim**: 32 (multi-embodiment), horizon 24
- **Training data**: MimicGen expert demos on LIBERO MimicGen-core (12 tasks)
- **Step**: 80,000

### `mimicgen-core-14b-full-step46000/`

- **Architecture**: same stock DreamZero, no changes
- **Variant**: full fine-tune (no LoRA) on 16 GPUs with DeepSpeed ZeRO
- **Sharding**: 10-shard safetensors (`model-{1..10}-of-00010.safetensors`)
- **Training data**: same MimicGen-core 12 tasks, longer instruction prompts
  ("detailed_instruct" recipe)
- **Step**: 46,000

## How to load

With the DreamZero codebase available:

```python
from stable_worldmodel.wm.utils import load_pretrained
# either subdir works the same way:
model = load_pretrained(
    "MinghaoFu/rlwf-ckpt/mimicgen-core-14b-lora-step80000",
    extra_args={"torch_dtype": "bfloat16"},
)
```

Direct safetensors load (LoRA, single file):

```python
from safetensors.torch import load_file
state_dict = load_file("model.safetensors")
```

Direct safetensors load (full, sharded):

```python
import json
from safetensors.torch import load_file

with open("model.safetensors.index.json") as f:
    index = json.load(f)
state_dict = {}
for shard in sorted(set(index["weight_map"].values())):
    state_dict.update(load_file(shard))
```

Full training config is in `experiment_cfg/conf.yaml` of each subdir.

## What is NOT in this repo

- DeepSpeed optimizer state (`global_step*/`) — stripped to keep the download
  small. If you want to resume training instead of just loading for inference,
  ping me; the optimizer shards are kept separately.
- `rng_state_*.pth` — same reason.
- The `latest` text file — points to a path inside `global_step*/`, irrelevant
  without the optimizer state.

## License

MIT (see `LICENSE`). The underlying Wan2.1-I2V-14B-480P base model has its own
Apache-2.0 license. DreamZero architecture follows the original
authors' release terms; this repo only redistributes the fine-tuned weights.

## Contact

Minghao Fu — isminghaofu@gmail.com