{"id":29551412,"url":"https://github.com/thuml/trajworld","last_synced_at":"2025-07-18T04:06:03.719Z","repository":{"id":297568823,"uuid":"990368562","full_name":"thuml/TrajWorld","owner":"thuml","description":"Official repository for \"Trajectory World Models for Heterogeneous Environments\" (ICML 2025), https://arxiv.org/abs/2502.01366","archived":false,"fork":false,"pushed_at":"2025-06-06T05:59:11.000Z","size":1209,"stargazers_count":3,"open_issues_count":0,"forks_count":0,"subscribers_count":1,"default_branch":"main","last_synced_at":"2025-06-06T06:30:21.548Z","etag":null,"topics":["model-predictive-control","off-policy-evaluation","pre-training","world-model"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":null,"status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/thuml.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":null,"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}},"created_at":"2025-05-26T02:36:07.000Z","updated_at":"2025-06-06T05:59:13.000Z","dependencies_parsed_at":"2025-06-06T06:41:35.881Z","dependency_job_id":null,"html_url":"https://github.com/thuml/TrajWorld","commit_stats":null,"previous_names":["thuml/trajworld"],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/thuml/TrajWorld","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/thuml%2FTrajWorld","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/thuml%2FTrajWorld/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/thuml%2FTrajWorld/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/thuml%2FTrajWorld/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/thuml","download_url":"https://codeload.github.com/thuml/TrajWorld/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/thuml%2FTrajWorld/sbom","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":265698049,"owners_count":23813128,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","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":["model-predictive-control","off-policy-evaluation","pre-training","world-model"],"created_at":"2025-07-18T04:05:54.576Z","updated_at":"2025-07-18T04:06:03.703Z","avatar_url":"https://github.com/thuml.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# Trajectory World Models for Heterogeneous Environments\n\n[![Paper](https://img.shields.io/badge/arXiv-Paper-b31b1b.svg?logo=arxiv)](https://arxiv.org/abs/2502.01366)\n\nThis is the official code base for the paper [Trajectory World Models for Heterogeneous Environments](https://arxiv.org/abs/2502.01366).\n\n![concept](assets/concept.png)\n\nGive it a star 🌟 if you find our work useful!\n\n## 🔥 News\n\n- 🚩 **2025.06.05**: Training code is released.\n- 🚩 **2025.05.01**: TrajWorld has been accepted by ICML 2025, congrats!.\n- 🚩 **2025.02.03**: Our paper is released on [arXiv](https://arxiv.org/abs/2502.01366).\n\n## 🛠️ Installation\n\n```bash\nconda env create -f environment.yaml\nconda activate trajworld\n```\n\n## 🌟 Pre-training\n\nTo pre-train the TrajWorld model:\n\n```bash\npython scripts/pretrain/pretrain_trajworld.py history_length=20 log_root_dir=log_pretrain_trajworld exp_name=merge_all n_blocks=6\n```\n\nSee more baseline scripts in [`scripts/pretrain`](scripts/pretrain).\n\n## 🎇 Fine-tuning\n\nTo fine-tune TrajWorld:\n\n```bash\npython scripts/training/train_trajworld.py env_name=hopper-medium-replay-v2 log_root_dir=log_model_new trm_epoch_steps=5000 dynamics_max_epochs_since_update=300 dynamics_max_epochs=50 seed=183 train_model_only=true exp_name=trajworld_ft trm_lr=1e-5 load_pt_dynamics_path=\"mergeall_pt/model/trm_dynamics990000.pkl\" n_blocks=6\n```\n\nTo train from scratch, simply remove the load_pt_dynamics_path argument.\n\nMore fine-tuning scripts for baseline models can be found in [`scripts/training`](scripts/training).\n\n## 📊 Evaluation\n\n### 🔁 Transition Prediction Evaluation\n\n```bash\npython pred/pred_mse_trajworld.py --env walker2d-random-v2 --model_path \u003cpath_to_your_model\u003e --n_blocks 6\n```\n\nAdditional baseline scripts:  [`pred`](pred).\n\n### 📈 OPE Evaluation\n\n```bash\npython ope/ope_eval.py --algo trajworld --env halfcheetah-expert-v2 --clear_kv_cache_every 10 --trm_lookback_window 10 --group 0 --n_blocks 6\n```\n\nSpecify your model path by modifying `get_list_dirs()` in [`ope/ope_eval.py`](ope/ope_eval.py).\n\nExample commands for other baselines are provided at the top of the file.\n\n### 🤖 MPC Evaluation\n\n```bash\npython mpc/mpc.py --algo trajworld --env walker2d-medium-replay-v2 --group 5 --clear_kv_cache_every 10 --trm_lookback_window 10 --action_proposal_id 3 --std 0.1\n```\n\nAs above, modify `get_list_dirs()` in  [`ope/ope_eval.py`](ope/ope_eval.py) to set your model path.\n\nBaseline examples are included in the script.\n\n\n## 🚀 Release Progress\n\n- [ ] UniTraj Dataset\n- [x] Pre-trained TrajWorld Model\n- [x] Transition prediction evaluation\n- [x] OPE evaluation\n- [x] MPC evaluation\n- [x] Training code for Trajworld, TDM and MLP-Ensemble\n\n## 📜 Citation\n\nIf you find this project useful, please cite our paper as:\n\n```\n@article{yin2025trajectory,\n  title={Trajectory World Models for Heterogeneous Environments},\n  author={Yin, Shaofeng and Wu, Jialong and Huang, Siqiao and Su, Xingjian and He, Xu and Hao, Jianye and Long, Mingsheng},\n  journal={arXiv preprint arXiv:2502.01366},\n  year={2025}\n}\n```\n\n## 🤝 Contact\n\nIf you have any questions, please contact yinshaofeng04@gmail.com.\n\n## 💡 Acknowledgement\n\nWe sincerely appreciate the following github repos for their valuable codebase we build upon:\n\n- https://github.com/nissymori/JAX-CORL\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fthuml%2Ftrajworld","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fthuml%2Ftrajworld","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fthuml%2Ftrajworld/lists"}