{"id":51551645,"url":"https://github.com/eljandoubi/driving-world-model","last_synced_at":"2026-07-10T00:30:52.911Z","repository":{"id":358309092,"uuid":"1240515311","full_name":"eljandoubi/driving-world-model","owner":"eljandoubi","description":"A world model for autonomous driving that predicts next-frame observations conditioned on actions, built on a UNet2D diffusion backbone with learned action embeddings.","archived":false,"fork":false,"pushed_at":"2026-05-23T20:15:52.000Z","size":264,"stargazers_count":0,"open_issues_count":0,"forks_count":0,"subscribers_count":0,"default_branch":"main","last_synced_at":"2026-05-23T22:13:09.104Z","etag":null,"topics":["diffusion-models","pytorch","unet-models","world-models"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/eljandoubi.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null,"notice":null,"maintainers":null,"copyright":null,"agents":null,"dco":null,"cla":null}},"created_at":"2026-05-16T08:18:10.000Z","updated_at":"2026-05-23T20:15:56.000Z","dependencies_parsed_at":null,"dependency_job_id":null,"html_url":"https://github.com/eljandoubi/driving-world-model","commit_stats":null,"previous_names":["eljandoubi/driving-world-model"],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/eljandoubi/driving-world-model","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/eljandoubi%2Fdriving-world-model","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/eljandoubi%2Fdriving-world-model/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/eljandoubi%2Fdriving-world-model/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/eljandoubi%2Fdriving-world-model/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/eljandoubi","download_url":"https://codeload.github.com/eljandoubi/driving-world-model/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/eljandoubi%2Fdriving-world-model/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":35316840,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-05-26T15:22:16.424Z","status":"online","status_checked_at":"2026-07-09T02:00:07.329Z","response_time":57,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["diffusion-models","pytorch","unet-models","world-models"],"created_at":"2026-07-10T00:30:51.785Z","updated_at":"2026-07-10T00:30:52.887Z","avatar_url":"https://github.com/eljandoubi.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# Driving World Model\n\nA world model for autonomous driving that learns to predict the next visual observation given the current frame and a driving action. The model is designed for hyperparameter tuning — all training parameters are exposed as CLI flags, making it straightforward to run sweeps with tools like W\u0026B Sweeps, Optuna, or grid search scripts.\n\n## Motivation\n\nWorld models enable planning and simulation without interacting with the real environment. This project frames next-frame prediction as a conditional generation task: given the current camera frame $x_t$ and a 3-dimensional action $a_t = (\\text{throttle}, \\text{steer}, \\text{brake})$, the model predicts $x_{t+1}$.\n\n## Architecture\n\n| Component | Description |\n|-----------|-------------|\n| **WorldModel** | Pretrained Stable Diffusion UNet conditioned on actions via cross-attention |\n| **ActionEmbedder** | Feed-forward network (SwiGLU) projecting $\\mathbb{R}^3 \\to \\mathbb{R}^{T \\times D}$ multi-token embeddings |\n| **Loss** | L1 (default) or L2 between predicted and ground-truth residual frame |\n\nThe UNet backbone is loaded from `runwayml/stable-diffusion-v1-5` (configurable). Actions are embedded into `num_tokens` vectors of dimension `cross_attention_dim` and injected as the encoder hidden states.\n\n## Dataset\n\nUses the [CARLA Autopilot Multimodal Dataset](https://huggingface.co/datasets/immanuelpeter/carla-autopilot-multimodal-dataset) streamed via HuggingFace `datasets`. Each sample pairs consecutive frames within the same driving run: $(x_t, a_t) \\to x_{t+1}$.\n\n## Installation\n\n```bash\nuv sync            # production deps\nuv sync --group dev  # + pytest for testing\n```\n\n## Training\n\n```bash\nuv run python train.py [OPTIONS]\n```\n\nAll fields in `TrainingConfig` are valid CLI flags (via HuggingFace `HfArgumentParser`). This makes it trivial to run hyperparameter sweeps:\n\n```bash\n# Example: sweep over learning rate and scheduler\nuv run python train.py --learning_rate 3e-4 --scheduler_t0 5 --scheduler_t_mult 1\nuv run python train.py --learning_rate 1e-4 --scheduler_t0 20 --scheduler_t_mult 2\n```\n\n### Resume from checkpoint\n\n```bash\nuv run python train.py --resume runs/\u003crun_id\u003e/checkpoints/best_checkpoint.pt --run_id \u003crun_id\u003e\n```\n\n### Distributed Training\n\nTraining supports DistributedDataParallel (DDP) for multi-GPU and multi-node setups. Data is sharded across ranks, gradients are synchronized automatically, and only rank 0 handles logging, checkpointing, and video generation.\n\n#### Single-node multi-GPU\n\n```bash\nuv run python train.py --n_gpus 4\n```\n\nThis uses `mp.spawn` to launch one process per GPU on the local machine.\n\n#### Multi-node\n\nEach machine runs the script independently. Set `MASTER_ADDR` to the IP of node 0 so all nodes can rendezvous.\n\n```bash\n# Node 0 (master):\nMASTER_ADDR=10.0.0.1 MASTER_PORT=12355 uv run python train.py \\\n    --n_gpus 4 --n_nodes 2 --node_rank 0\n\n# Node 1:\nMASTER_ADDR=10.0.0.1 MASTER_PORT=12355 uv run python train.py \\\n    --n_gpus 4 --n_nodes 2 --node_rank 1\n```\n\n#### Multi-node with `torchrun`\n\n`torchrun` handles rendezvous and environment variables automatically. When `LOCAL_RANK` is detected in the environment, the script reads rank/world-size from `torchrun` instead of the CLI flags.\n\n```bash\n# Node 0:\ntorchrun --nproc_per_node=4 --nnodes=2 --node_rank=0 \\\n    --master_addr=10.0.0.1 --master_port=12355 train.py\n\n# Node 1:\ntorchrun --nproc_per_node=4 --nnodes=2 --node_rank=1 \\\n    --master_addr=10.0.0.1 --master_port=12355 train.py\n```\n\n## Default Configuration\n\n### Model\n\n| Parameter | Default | Description |\n|-----------|---------|-------------|\n| `base_name` | `runwayml/stable-diffusion-v1-5` | Pretrained UNet backbone |\n| `num_tokens` | `16` | Number of action embedding tokens |\n| `hidden_dim` | `2048` | Action embedder hidden dimension |\n| `activation` | `swiglu` | Activation in action embedder (`gelu`, `geglu`, `swiglu`, etc.) |\n| `loss_type` | `l1` | Loss function (`l1` or `l2`) |\n| `dropout` | `0.0` | Dropout in action embedder |\n\n### Optimization\n\n| Parameter | Default | Description |\n|-----------|---------|-------------|\n| `learning_rate` | `1e-4` | AdamW learning rate |\n| `max_grad_norm` | `1.0` | Gradient clipping max norm |\n| `scheduler_t0` | `10` | CosineAnnealingWarmRestarts $T_0$ (epochs per first restart) |\n| `scheduler_t_mult` | `2` | CosineAnnealingWarmRestarts $T_{mult}$ (period multiplier) |\n| `epochs` | `100000` | Maximum training epochs |\n\n### Early Stopping \u0026 Checkpointing\n\n| Parameter | Default | Description |\n|-----------|---------|-------------|\n| `patience` | `10` | Early stopping patience (number of validations without improvement) |\n| `min_delta` | `1e-8` | Minimum improvement to reset patience counter |\n| `log_every` | `1000` | Log training loss every N steps |\n| `checkpoint_every` | `10000` | Validate \u0026 checkpoint every N steps |\n\n### Distributed\n\n| Parameter | Default | Description |\n|-----------|---------|-------------|\n| `n_gpus` | `1` | GPUs per node |\n| `n_nodes` | `1` | Number of nodes (machines) |\n\n### Data\n\n| Parameter | Default | Description |\n|-----------|---------|-------------|\n| `image_size` | `256` | Input image resize resolution |\n| `batch_size` | `2` | Training batch size (per GPU) |\n| `buffer_size` | `1000` | Prefetch buffer size for streaming dataset |\n| `pin_memory` | `True` | Pin memory for GPU transfer |\n\n### Paths\n\n| Parameter | Default | Description |\n|-----------|---------|-------------|\n| `runs_dir` | `runs` | Base directory for all run outputs |\n| `checkpoint_dir` | `checkpoints` | Subdirectory for checkpoints (under `runs/\u003crun_id\u003e/`) |\n| `plot_dir` | `plots` | Subdirectory for plots (under `runs/\u003crun_id\u003e/`) |\n| `resume` | `\"\"` | Path to checkpoint to resume from |\n| `run_id` | `None` | W\u0026B run ID (auto-generated if not set) |\n\n## Hyperparameter Tuning\n\nKey parameters to sweep:\n\n- **`learning_rate`**: Try log-uniform in `[1e-5, 1e-3]`\n- **`num_tokens`**: `{4, 8, 16}` — controls action conditioning capacity\n- **`hidden_dim`**: `{512, 1024, 2048}` — action embedder width\n- **`scheduler_t0`** / **`scheduler_t_mult`**: Controls warm restart frequency\n- **`dropout`**: `[0.0, 0.3]`\n- **`image_size`**: `{128, 256, 512}` — resolution/compute trade-off\n- **`batch_size`**: `{4, 8, 16, 32}` — adjust with available GPU memory\n\nAll parameters are logged to W\u0026B automatically, enabling comparison across runs.\n\n## Testing\n\n```bash\nuv run pytest tests/ -v\n```\n\n## Project Structure\n\n```\n├── model.py            # WorldModel + ActionEmbedder\n├── train.py            # Training loop with scheduler, early stopping, checkpointing\n├── validate.py         # Validation/test evaluation\n├── dataset.py          # Streaming dataset from HuggingFace\n├── config.py           # TrainingConfig dataclass (all hyperparameters)\n├── checkpoint.py       # Save/load checkpoints\n├── early_stopping.py   # EarlyStopping utility\n├── logger.py           # Rank-aware logging with tqdm integration\n├── plot.py             # Video generation for prediction visualization\n├── tests/              # Unit tests\n└── pyproject.toml      # Dependencies \u0026 config\n```\n\n## License\n\nSee [LICENSE](LICENSE).","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Feljandoubi%2Fdriving-world-model","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Feljandoubi%2Fdriving-world-model","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Feljandoubi%2Fdriving-world-model/lists"}