Skip to content

feature(xjy): Refine PriorZero Implementation #441

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
xiongjyu wants to merge 187 commits into
base: dev-multitask-balance-clean-rft
Choose a base branch
from
Open

feature(xjy): Refine PriorZero Implementation #441

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
187 commits
Select commit Hold shift + click to select a range
a3a2d69
Fix game_segment/weighted_total_loss bugs and refine prompts, compute…
xiongjyu Nov 20, 2025
959a558
Fixed the accumulate_steps bug and added cprofile functionality.
xiongjyu Nov 20, 2025
ecedc5f
Refine the code and fix the bug in data collection.
xiongjyu Nov 22, 2025
2d53d22
Add REINFORCE-style losses and store old_logprob in the buffer.
xiongjyu Nov 23, 2025
c608600
Fix the get_llm_prior bug so that every action receives a logprob
xiongjyu Nov 24, 2025
15e39f6
fixed the history bug in the build_llm_prompt and logs in forward_learn
xiongjyu Nov 24, 2025
7c9acd9
rename advantage_tensor on rft
xiongjyu Nov 24, 2025
738f300
Fixed the action out-of-bounds bug and added a record for forward_col…
xiongjyu Nov 26, 2025
0a166f6
Fixed the misalignment between old_log_prob and log_prob, and correct…
xiongjyu Nov 27, 2025
4f3668e
add some logs for analysying
xiongjyu Nov 27, 2025
2985e60
Polish the code and standardize the format.
xiongjyu Nov 29, 2025
ff98006
Add kL divergence in rft and llm_prior_entropy in collect
xiongjyu Dec 2, 2025
7e43e45
polish config and format
xiongjyu Dec 3, 2025
d6555e5
delete unused files
xiongjyu Dec 3, 2025
b7d42ee
Decouple the training of world_model and LLM.
xiongjyu Dec 9, 2025
95e2347
add cache in the jericho
xiongjyu Dec 10, 2025
9682486
Separate sync and async entry points to simplify the program.
xiongjyu Dec 10, 2025
0a38197
Reference OpenRLHF’s implementation to update vLLM weights in real ti…
xiongjyu Dec 14, 2025
e361039
delete unused orz files
xiongjyu Dec 14, 2025
f957db9
fix a small bug
xiongjyu Dec 15, 2025
628d7d2
Fix action='go' bug; optimize replay buffer with larger capacity; sam…
xiongjyu Dec 16, 2025
c16174f
fix a bug
xiongjyu Dec 16, 2025
35cb4f9
Optimized log-probability computation for the CoT setting.
xiongjyu Dec 17, 2025
2c67a8d
polish and format file
xiongjyu Dec 19, 2025
b16c3e7
Improve single/multi-process LLM training with DeepSpeed
xiongjyu Dec 26, 2025
97c9843
fix the vllm bug when using torchrun
xiongjyu Dec 27, 2025
eb8a4bd
fix the fork bug and polish the vllm about sleep
xiongjyu Dec 27, 2025
9178397
Optimized efficiency and added multiple ways to calculate advantage.
xiongjyu Dec 28, 2025
cb6f7cf
limit the ouput length when using vllm
xiongjyu Dec 28, 2025
19fac8f
polish(pu): add cot-reuse in training, use running-norm in value, pol…
puyuan1996 Dec 29, 2025
88f047b
fix(pu): fix some bugs in reuse-collect-cot in training phase
puyuan1996 Dec 29, 2025
de4b2c0
polish configs and format
xiongjyu Dec 30, 2025
2069e32
delete unuse config
xiongjyu Dec 30, 2025
3ff091e
fix not found go bug
xiongjyu Dec 30, 2025
3888d8e
fix the misalignment bug when reusing cot
xiongjyu Jan 6, 2026
da0d0fd
make the prompt more compact
xiongjyu Jan 6, 2026
5f88151
add lr warmup
xiongjyu Jan 7, 2026
ed89062
add warmup for training world model before training llm and AdaptiveV…
xiongjyu Jan 7, 2026
96dc250
polish the implementation of profile
xiongjyu Jan 7, 2026
66ac376
fix a small bug
xiongjyu Jan 8, 2026
a9593cd
add profile of forward_collect
xiongjyu Jan 8, 2026
5d0f359
add format reward option and fix the cot gradient
xiongjyu Jan 8, 2026
e15e66c
rename kl/clip-ratio metrics
xiongjyu Jan 14, 2026
55e66ed
Optimize the use of format rewards
xiongjyu Jan 14, 2026
0cef8b8
polish the format
xiongjyu Jan 14, 2026
f88989b
Add a complete DDP program, including the process of collect and worl…
xiongjyu Jan 14, 2026
ffcf348
fix some small bugs
xiongjyu Jan 14, 2026
591c420
add the result of valid_actions's prob in vllm_output and samples shu…
xiongjyu Jan 15, 2026
091cc80
add the policy model for reload/offload func
xiongjyu Jan 16, 2026
e512707
fix a bug
xiongjyu Jan 16, 2026
a6bdc0e
Fixed a bug in the advantage feature and used target_value-pred_value…
xiongjyu Jan 17, 2026
250ba63
Optimize parameter definitions and off-policy's implementation
xiongjyu Jan 18, 2026
c3af1c2
fix a small bug
xiongjyu Jan 18, 2026
9d304c6
fix a small bug
xiongjyu Jan 19, 2026
9654f6a
Optimize essential logging; add grad norm metrics; record metrics as …
xiongjyu Feb 1, 2026
1bbcc19
fix a small bug
xiongjyu Feb 2, 2026
335e16c
Fix a bug; refactor LLM prompts into system/user roles; improve CoT o…
xiongjyu Feb 2, 2026
acf5d04
add input/response length metrics and envstep-based tb_logger for lea…
xiongjyu Feb 4, 2026
2ff4a90
add llm_prior_tempearture to apply_temperature_scaling for llm_prior
xiongjyu Feb 5, 2026
ea32a4a
Fixed bugs in prefix_cot & padding and aligned the input contexts of …
xiongjyu Feb 7, 2026
f895ed2
Remove the logits_to_keep parameter and add metrics such as entropy.
xiongjyu Feb 8, 2026
4567190
tmp
xiongjyu Feb 8, 2026
5098395
refine some log
xiongjyu Feb 9, 2026
a1f8282
feature(pu): add init image-based/vlm version of priorzero
puyuan1996 Feb 11, 2026
84b3317
fix(pu): fix vlm_engine.py
puyuan1996 Feb 11, 2026
b865dcb
polish some config
xiongjyu Feb 23, 2026
82e1d29
Fixed the bug caused by fmt_weight being 1
xiongjyu Feb 24, 2026
7c9c922
delete unused file
xiongjyu Feb 24, 2026
a13292a
tmp
puyuan1996 Feb 24, 2026
2dcc4ae
fix score, action_str, and timestep_dict bug, and add 3 modes of Prio…
xiongjyu Feb 25, 2026
0cf4bb7
fix(pu): fix tb log in priorzero_evaluator.py
puyuan1996 Feb 25, 2026
a8af654
Merge branch 'dev-multitask-balance-clean-rft-from-0225' into dev-mul…
puyuan1996 Feb 25, 2026
d73c16c
fix(pu): fix vision token in prompts in vlm settings
puyuan1996 Feb 25, 2026
d89a3c5
feature(pu): adapt prior_generator.py to support priorzero training i…
puyuan1996 Feb 25, 2026
f70d6d7
polish(pu): polish logs
puyuan1996 Feb 25, 2026
a31b5a1
polish(pu): polish logs
puyuan1996 Feb 25, 2026
276761b
test(pu): add priorzero pong test config
puyuan1996 Feb 25, 2026
dbec27c
fix self.history_buffer bug and refine logs
xiongjyu Feb 26, 2026
e71eb61
fix some bug when running unizero
xiongjyu Feb 28, 2026
e3f7cfd
fix a bug in evaluator and add enable_rft/wm to control what models n…
xiongjyu Feb 28, 2026
7377220
add bash scripts to run priorzero and format files
xiongjyu Feb 28, 2026
d5c4923
tmp
xiongjyu Feb 28, 2026
3caf28f
add priorzero README.md
xiongjyu Feb 28, 2026
cfc000a
refine cot prefix and add mcts_root_logits_dict
xiongjyu Feb 28, 2026
8bec5f7
tmp
xiongjyu Feb 28, 2026
f49c9d9
add user_prompt_dict to control user_prompt about reward/valid_action…
xiongjyu Feb 28, 2026
1f0327e
Merge remote-tracking branch 'origin-xjy/dev-multitask-balance-clean-…
puyuan1996 Mar 2, 2026
9a459f8
fix the bug of pretrained_model and some small bugs
xiongjyu Mar 4, 2026
38efe3d
fix the evaluator_env_num when steping into the _init_eval
xiongjyu Mar 4, 2026
0002c8b
Add configuration and adaptation code for alternating LLM and WM trai…
xiongjyu Mar 5, 2026
881d9c4
adapter the wm to load pretrained weight
xiongjyu Mar 5, 2026
1955f6b
rename the log_name in scripts
xiongjyu Mar 5, 2026
31f2090
polish alternate configuration and add eval_episode_log
xiongjyu Mar 6, 2026
ce44f7d
fix the env_id bug in mcts_info
xiongjyu Mar 6, 2026
c87f0f1
sync global_step in priorzero_trainer
xiongjyu Mar 7, 2026
cf1d1cf
add adaptive mode for llm/wm policy
xiongjyu Mar 8, 2026
a4073be
fix the bug that llm_weight could less 0
xiongjyu Mar 10, 2026
c5c9f44
merge pr-474: reset parent_q per batch element in batch_traverse
xiongjyu Mar 10, 2026
126228b
output llm_weight in collect log and optimizer the way of llm_plus_wm…
xiongjyu Mar 10, 2026
9be32c3
add llm_sft code
xiongjyu Mar 11, 2026
8ffb05b
format _collect/eval suffix in the unizero
xiongjyu Mar 12, 2026
eb27b00
tmp
xiongjyu Mar 12, 2026
4c77e6c
Fixed the bug where use_cot was False and added a function to clear v…
xiongjyu Mar 12, 2026
8328bbb
fix llm_weight bug in priorzero_policy and polish run_priorzero_ddp.sh
xiongjyu Mar 12, 2026
e10dd30
Fixed a deadlock bug that occurred when saving LLM weights while runn…
xiongjyu Mar 13, 2026
3c9e232
add lora and statistics on data duplication rate
xiongjyu Mar 14, 2026
68d05cb
set the advantage_batch_norm and llm_plus_wm_logits to default
xiongjyu Mar 15, 2026
b80df0c
tmp
xiongjyu Mar 15, 2026
08221be
fix a bug in priorzero_collector for env_id misalignment and a bug i…
xiongjyu Mar 18, 2026
ede1637
Optimize the data sampling process for LLM training
xiongjyu Mar 18, 2026
0421f45
fix the bug in run_ddp
xiongjyu Mar 18, 2026
8731968
Several bugs were fixed and vllm parameters were adjusted.
xiongjyu Mar 19, 2026
b946106
Add `old_action_logprob` and `rollout_action_logprob` to fix the bug …
xiongjyu Mar 19, 2026
715c776
fix a small bug to prevent OOM
xiongjyu Mar 20, 2026
490dd88
optimizer priorzero_entry_ddp and add unique samples for make_llm_sam…
xiongjyu Mar 20, 2026
823f9cd
Add torch.cuda.empty_cache() after fetch_latest_batch to prevent OOM
xiongjyu Mar 20, 2026
38dc3cf
fix a small bug
xiongjyu Mar 20, 2026
71d2720
Merge remote-tracking branch 'origin-xjy/dev-multitask-balance-clean-…
puyuan1996 Mar 20, 2026
e51f5fd
feature/fix(pu): fix priorzero-vl, add init version of run_priorzero_…
puyuan1996 Mar 20, 2026
9436fd5
polish(pu): rename vlm to vl, polish file structure
puyuan1996 Mar 20, 2026
fcabce5
fix(pu): fix some bugs in pipeline of run_priorzero_vl_lunarlander.sh
Mar 20, 2026
7283e47
polish(pu): optimize vlm prompt and config
puyuan1996 Mar 20, 2026
6ea50c8
polish(pu): polish eval and config in run_priorzero_vl_lunarlander.sh…
puyuan1996 Mar 21, 2026
f5db0db
polish(pu): polish prompts and args
puyuan1996 Mar 21, 2026
f04bab5
Expand the range of pos_in_game_segment and add gradient weights for …
xiongjyu Mar 21, 2026
e641277
fix(pu): fix sampling-params, fix add-sample to buffer, polish logs
puyuan1996 Mar 21, 2026
3b213c7
fix(pu): fix raw_obs_text name bug in priorzero_collector_unified.py,…
puyuan1996 Mar 21, 2026
d0283c5
for llm train, collector only to collect data without mcts+wm
xiongjyu Mar 21, 2026
318c6be
prevent OOM when running ddp
xiongjyu Mar 21, 2026
94f5b00
fix(pu): fix tb, polish prompt, fix timestep in game_history
puyuan1996 Mar 21, 2026
56878c2
fix a bug when arriving llm phase and collect process
xiongjyu Mar 21, 2026
b3c9f45
polish(pu): add some monitor metrics in tb
puyuan1996 Mar 21, 2026
59c7972
remove the pure-llm to collect, add pure-wm to collect when running l…
xiongjyu Mar 22, 2026
e6bf3a6
rename exp_name
xiongjyu Mar 22, 2026
c80cab3
fix(pu): fix lunarlander prompts, fix vl_engine
puyuan1996 Mar 22, 2026
036671a
polish(pu): add vlm_image_mode option in lunarlander-image priorzero
puyuan1996 Mar 22, 2026
1371ef5
fix some inportant bugs to prevent Nan
xiongjyu Mar 23, 2026
e2c37af
feature(pu): add lunarlander_image_unizero_config.py
puyuan1996 Mar 23, 2026
0a13c6a
tmp
xiongjyu Mar 23, 2026
43d322a
polish the process of llm's training samples
xiongjyu Mar 23, 2026
dee830d
polish(pu): polish lunarlander config
puyuan1996 Mar 24, 2026
a7daccc
Merge remote-tracking branch 'origin-xjy/dev-multitask-balance-clean-…
puyuan1996 Mar 24, 2026
a5411bd
feature(pu): add cot_weight option in priorzero-vl
puyuan1996 Mar 24, 2026
1c11f61
fix(pu): mcts-action bug in priorzero-vl, add eval_vl_prior and run_p…
puyuan1996 Mar 24, 2026
f07a175
fix(pu): fix return in muzero_evaluator.py for lunarlander_image_uni…
puyuan1996 Mar 24, 2026
95a5bd9
polish evaluator log and eval_freq
xiongjyu Mar 25, 2026
331a23d
tmp
xiongjyu Mar 25, 2026
059cbbc
add advantage_global_batch_norm
xiongjyu Mar 25, 2026
be5f941
tmp
xiongjyu Mar 25, 2026
706da1b
delete unused files and configs
xiongjyu Mar 27, 2026
907ab05
Merge remote-tracking branch 'origin-xjy/dev-multitask-balance-clean-…
puyuan1996 Mar 28, 2026
703797d
fix(pu): polish lunarlander_image_unizero_config, fix logprob in prio…
puyuan1996 Mar 28, 2026
fa11cd7
fix(pu): add logprob_extraction_mode to support two modes: exact and …
puyuan1996 Mar 28, 2026
9159a0e
add valid_actions to prompt
xiongjyu Mar 30, 2026
14c09ca
fix a bug when valid_actions is greater max_action_nums in evaluating…
xiongjyu Apr 1, 2026
29a74c9
fix a bug when zork1 encounter the emulator halted
xiongjyu Apr 2, 2026
b515254
fix zork1 stucking the step() using timeout judgment
xiongjyu Apr 4, 2026
d040b93
fix a bug
xiongjyu Apr 4, 2026
9095783
fix timeout bug when running zork1
xiongjyu Apr 4, 2026
09ddd49
tmp
xiongjyu Apr 5, 2026
a6ae249
fix(pu): fix jericho c timeout bug
puyuan1996 Apr 5, 2026
888ae2a
fix the bug duing to zork1
xiongjyu Apr 9, 2026
caed1c1
add llm_no_collect mode
xiongjyu Apr 9, 2026
6f30b73
docs(priorzero): 补充Atari等vision环境改进方向
xiongjyu Apr 11, 2026
c575dfc
Merge pull request #10 from xiongjyu/codex/-zoo/jericho/priorzero-readme
xiongjyu Apr 11, 2026
0b93856
delete unused files
xiongjyu Apr 11, 2026
40c7d48
Merge remote-tracking branch 'origin-xjy/dev-multitask-balance-clean-…
puyuan1996 Apr 22, 2026
4d50c4d
polish(pu): polish lunarlander-image uz config
puyuan1996 Apr 22, 2026
02ea915
polish(pu): polish lunarlander-image related code
puyuan1996 Apr 22, 2026
dad9df9
tmp
xiongjyu Apr 23, 2026
e318dcc
feature(pu): add babyai env and related priorzero configs
puyuan1996 Apr 25, 2026
b51efff
fix(pu): fix babyai return
puyuan1996 Apr 25, 2026
58a9824
feature(pu): add textcraft env and configs
puyuan1996 Apr 25, 2026
68ae363
add llm as policy and rlft ablations
xiongjyu Apr 28, 2026
1a59d6a
polish(pu): Align BabyAI config with ScalingInter-RL: multi-task 40 l…
puyuan1996 Apr 28, 2026
f0c236b
fix(pu): Fix BERT 1D tensor crash in eval, align BabyAI to 18 Scaling…
puyuan1996 Apr 28, 2026
6793702
add qwen as backbone
xiongjyu Apr 28, 2026
96001b2
fix(pu): TP-pair drain in collector/eval, structured logging, eval-tr…
puyuan1996 Apr 28, 2026
2309ab6
tmp
xiongjyu Apr 28, 2026
fe14fab
fix(pu): Fix CUDA OOM in chunked forward by removing hardcoded chunk_…
puyuan1996 Apr 29, 2026
fb53d67
fix(pu): Fix LLM eval coverage, stabilize RFT training, add UniZero b…
puyuan1996 Apr 30, 2026
41a212b
fix(pu): Unify eval x-axis to env_step, add per-level UniZero evaluat…
puyuan1996 May 2, 2026
35d4c67
fix(pu): fix babyai_env.py
puyuan1996 May 12, 2026
e5648f9
Merge branch 'dev-multitask-balance-clean-rft' into dev-multitask-bal…
xiongjyu Jun 13, 2026
d3b2b60
Merge pull request #11 from opendilab/dev-multitask-balance-clean-rft…
xiongjyu Jun 13, 2026
c986457
Merge branch 'dev-multitask-balance-clean-rft' into dev-multitask-bal…
xiongjyu Jun 15, 2026
e17d3ee
Merge pull request #8 from opendilab/dev-multitask-balance-clean-rft-vl
xiongjyu Jun 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion .gitignore
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,10 @@ data_*
*.gv
*.png
*.csv
*.jsonl
*.json

pkg/
src/

### CVS template
/CVS/*
Expand Down
Loading