feature(xjy): Refine PriorZero Implementation - #441

Open

xiongjyu wants to merge 187 commits into

opendilab:dev-multitask-balance-clean-rftfrom

xiongjyu:dev-multitask-balance-clean-rft

Commits on Nov 20, 2025

Fix game_segment/weighted_total_loss bugs and refine prompts, compute_llm_prior, and SFT loss
xiongjyu
committed
a3a2d69
Fixed the accumulate_steps bug and added cprofile functionality.
xiongjyu
committed
959a558

Commits on Nov 22, 2025

Refine the code and fix the bug in data collection.
xiongjyu
committed
ecedc5f

Commits on Nov 23, 2025

Add REINFORCE-style losses and store old_logprob in the buffer.
xiongjyu
committed
2d53d22

Commits on Nov 24, 2025

Fix the get_llm_prior bug so that every action receives a logprob
xiongjyu
committed
c608600
fixed the history bug in the build_llm_prompt and logs in forward_learn
xiongjyu
committed
15e39f6
rename advantage_tensor on rft
xiongjyu
committed
7c9acd9

Commits on Nov 26, 2025

Fixed the action out-of-bounds bug and added a record for forward_collect to cprofile.
xiongjyu
committed
738f300

Commits on Nov 27, 2025

Commits on Nov 29, 2025

Polish the code and standardize the format.
xiongjyu
committed
2985e60

Commits on Dec 2, 2025

Add kL divergence in rft and llm_prior_entropy in collect
xiongjyu
committed
ff98006

Commits on Dec 3, 2025

polish config and format
xiongjyu
committed
7e43e45
delete unused files
xiongjyu
committed
d6555e5

Commits on Dec 9, 2025

Decouple the training of world_model and LLM.
xiongjyu
committed
b7d42ee

Commits on Dec 10, 2025

add cache in the jericho
xiongjyu
committed
95e2347
Separate sync and async entry points to simplify the program.
xiongjyu
committed
9682486

Commits on Dec 14, 2025

Commits on Dec 15, 2025

fix a small bug
xiongjyu
committed
f957db9

Commits on Dec 16, 2025

Commits on Dec 17, 2025

Optimized log-probability computation for the CoT setting.
xiongjyu
committed
35cb4f9

Commits on Dec 19, 2025

polish and format file
xiongjyu
committed
2c67a8d

Commits on Dec 26, 2025

Improve single/multi-process LLM training with DeepSpeed
xiongjyu
committed
b16c3e7

Commits on Dec 27, 2025

fix the vllm bug when using torchrun
xiongjyu
committed
97c9843
fix the fork bug and polish the vllm about sleep
xiongjyu
committed
eb8a4bd

Commits on Dec 28, 2025

Optimized efficiency and added multiple ways to calculate advantage.
xiongjyu
committed
9178397
limit the ouput length when using vllm
xiongjyu
committed
cb6f7cf

Commits on Dec 29, 2025

Commits on Dec 30, 2025

polish configs and format
xiongjyu
committed
de4b2c0
delete unuse config
xiongjyu
committed
2069e32
fix not found go bug
xiongjyu
committed
3ff091e

Commits on Jan 6, 2026

fix the misalignment bug when reusing cot
xiongjyu
committed
3888d8e
make the prompt more compact
xiongjyu
committed
da0d0fd

Commits on Jan 7, 2026

add lr warmup
xiongjyu
committed
5f88151
add warmup for training world model before training llm and AdaptiveValueNormalizer
xiongjyu
committed
ed89062
polish the implementation of profile
xiongjyu
committed
96dc250

Commits on Jan 8, 2026

fix a small bug
xiongjyu
committed
66ac376
add profile of forward_collect
xiongjyu
committed
a9593cd
add format reward option and fix the cot gradient
xiongjyu
committed
5d0f359

Commits on Jan 14, 2026

Commits on Jan 15, 2026

add the result of valid_actions's prob in vllm_output and samples shuffle and some fine-grained metrics
xiongjyu
committed
591c420

Commits on Jan 16, 2026

add the policy model for reload/offload func
xiongjyu
committed
091cc80
fix a bug
xiongjyu
committed
e512707

Commits on Jan 17, 2026

Fixed a bug in the advantage feature and used target_value-pred_value as the advantage value.
xiongjyu
committed
a6bdc0e

Commits on Jan 18, 2026

Optimize parameter definitions and off-policy's implementation
xiongjyu
committed
250ba63
fix a small bug
xiongjyu
committed
c3af1c2

Commits on Jan 19, 2026

fix a small bug
xiongjyu
committed
9d304c6

Commits on Feb 1, 2026

Optimize essential logging; add grad norm metrics; record metrics as averages over parameter updates; add NaN debug logs; enable LLM checkpoint saving.
xiongjyu
committed
9654f6a

Commits on Feb 2, 2026

fix a small bug
xiongjyu
committed
1bbcc19
Fix a bug; refactor LLM prompts into system/user roles; improve CoT outputs.
xiongjyu
committed
335e16c

Commits on Feb 4, 2026

add input/response length metrics and envstep-based tb_logger for learner_llm
xiongjyu
committed
acf5d04

Commits on Feb 5, 2026

add llm_prior_tempearture to apply_temperature_scaling for llm_prior
xiongjyu
committed
2ff4a90

Commits on Feb 7, 2026

Fixed bugs in prefix_cot & padding and aligned the input contexts of policy_model and vllm.
xiongjyu
committed
ea32a4a

Commits on Feb 8, 2026

Remove the logits_to_keep parameter and add metrics such as entropy.
xiongjyu
committed
f895ed2
tmp
xiongjyu
committed
4567190

Commits on Feb 9, 2026

refine some log
xiongjyu
committed
5098395

Commits on Feb 11, 2026

feature(pu): add init image-based/vlm version of priorzero
puyuan1996
committed
a1f8282
fix(pu): fix vlm_engine.py
puyuan1996
committed
84b3317

Commits on Feb 23, 2026

polish some config
xiongjyu
committed
b865dcb

Commits on Feb 24, 2026

Fixed the bug caused by fmt_weight being 1
xiongjyu
committed
82e1d29
delete unused file
xiongjyu
committed
7c9c922
tmp
puyuan1996
committed
a13292a

Commits on Feb 25, 2026

Commits on Feb 26, 2026

fix self.history_buffer bug and refine logs
xiongjyu
committed
dbec27c

Commits on Feb 28, 2026

fix some bug when running unizero
xiongjyu
committed
e71eb61
fix a bug in evaluator and add enable_rft/wm to control what models need to be trained
xiongjyu
committed
e3f7cfd
add bash scripts to run priorzero and format files
xiongjyu
committed
7377220
tmp
xiongjyu
committed
d5c4923
add priorzero README.md
xiongjyu
committed
3caf28f
refine cot prefix and add mcts_root_logits_dict
xiongjyu
committed
cfc000a
tmp
xiongjyu
committed
8bec5f7
add user_prompt_dict to control user_prompt about reward/valid_action(valid_action not completed)
xiongjyu
committed
f49c9d9

Commits on Mar 2, 2026

Merge remote-tracking branch 'origin-xjy/dev-multitask-balance-clean-rft' into dev-multitask-balance-clean-rft-0210-vlm
puyuan1996
committed
1f0327e

Commits on Mar 4, 2026

fix the bug of pretrained_model and some small bugs
xiongjyu
committed
9a459f8
fix the evaluator_env_num when steping into the _init_eval
xiongjyu
committed
38efe3d

Commits on Mar 5, 2026

Add configuration and adaptation code for alternating LLM and WM training.
xiongjyu
committed
0002c8b
adapter the wm to load pretrained weight
xiongjyu
committed
881d9c4
rename the log_name in scripts
xiongjyu
committed
1955f6b

Commits on Mar 6, 2026

polish alternate configuration and add eval_episode_log
xiongjyu
committed
31f2090
fix the env_id bug in mcts_info
xiongjyu
committed
ce44f7d

Commits on Mar 7, 2026

sync global_step in priorzero_trainer
xiongjyu
committed
c87f0f1

Commits on Mar 8, 2026

add adaptive mode for llm/wm policy
xiongjyu
committed
cf1d1cf

Commits on Mar 10, 2026

Commits on Mar 11, 2026

add llm_sft code
xiongjyu
committed
9be32c3

Commits on Mar 12, 2026

Commits on Mar 13, 2026

Fixed a deadlock bug that occurred when saving LLM weights while running DDP files.
xiongjyu
committed
e10dd30

Commits on Mar 14, 2026

add lora and statistics on data duplication rate
xiongjyu
committed
3c9e232

Commits on Mar 15, 2026

set the advantage_batch_norm and llm_plus_wm_logits to default
xiongjyu
committed
68d05cb
tmp
xiongjyu
committed
b80df0c

Commits on Mar 18, 2026

fix a bug in priorzero_collector for env_id misalignment and a bug in last_pos_in_transition
xiongjyu
committed
08221be
Optimize the data sampling process for LLM training
xiongjyu
committed
ede1637
fix the bug in run_ddp
xiongjyu
committed
0421f45

Commits on Mar 19, 2026

Commits on Mar 20, 2026

fix a small bug to prevent OOM
xiongjyu
committed
715c776
optimizer priorzero_entry_ddp and add unique samples for make_llm_samples
xiongjyu
committed
490dd88
Add torch.cuda.empty_cache() after fetch_latest_batch to prevent OOM
xiongjyu
committed
823f9cd
fix a small bug
xiongjyu
committed
38dc3cf
Merge remote-tracking branch 'origin-xjy/dev-multitask-balance-clean-rft' into dev-multitask-balance-clean-rft-0210-vlm
puyuan1996
committed
71d2720
feature/fix(pu): fix priorzero-vl, add init version of run_priorzero_vlm_lunarlander.sh
puyuan1996
committed
e51f5fd
polish(pu): rename vlm to vl, polish file structure
puyuan1996
committed
9436fd5
fix(pu): fix some bugs in pipeline of run_priorzero_vl_lunarlander.sh

root
authored and
puyuan1996
committed
fcabce5
polish(pu): optimize vlm prompt and config
puyuan1996
committed
7283e47

Commits on Mar 21, 2026

Commits on Mar 22, 2026

Commits on Mar 23, 2026

fix some inportant bugs to prevent Nan
xiongjyu
committed
1371ef5
feature(pu): add lunarlander_image_unizero_config.py
puyuan1996
committed
e2c37af
tmp
xiongjyu
committed
0a13c6a
polish the process of llm's training samples
xiongjyu
committed
43d322a

Commits on Mar 24, 2026

Commits on Mar 25, 2026

polish evaluator log and eval_freq
xiongjyu
committed
95a5bd9
tmp
xiongjyu
committed
331a23d
add advantage_global_batch_norm
xiongjyu
committed
059cbbc
tmp
xiongjyu
committed
be5f941

Commits on Mar 27, 2026

delete unused files and configs
xiongjyu
committed
706da1b

Commits on Mar 28, 2026

Commits on Mar 30, 2026

add valid_actions to prompt
xiongjyu
committed
9159a0e

Commits on Apr 1, 2026

fix a bug when valid_actions is greater max_action_nums in evaluating wm_llm_prior
xiongjyu
committed
14c09ca

Commits on Apr 2, 2026

fix a bug when zork1 encounter the emulator halted
xiongjyu
committed
29a74c9

Commits on Apr 4, 2026

fix zork1 stucking the step() using timeout judgment
xiongjyu
committed
b515254
fix a bug
xiongjyu
committed
d040b93
fix timeout bug when running zork1
xiongjyu
committed
9095783

Commits on Apr 5, 2026

tmp
xiongjyu
committed
09ddd49
fix(pu): fix jericho c timeout bug
puyuan1996
committed
a6ae249

Commits on Apr 9, 2026

fix the bug duing to zork1
xiongjyu
committed
888ae2a
add llm_no_collect mode
xiongjyu
committed
caed1c1

Commits on Apr 11, 2026

docs(priorzero): 补充Atari等vision环境改进方向
 xiongjyu
committed
6f30b73
Merge pull request #10 from xiongjyu/codex/-zoo/jericho/priorzero-readme
xiongjyu
authored
c575dfc
delete unused files
xiongjyu
committed
0b93856

Commits on Apr 22, 2026

Commits on Apr 23, 2026

tmp
xiongjyu
committed
dad9df9

Commits on Apr 25, 2026

feature(pu): add babyai env and related priorzero configs
puyuan1996
committed
e318dcc
fix(pu): fix babyai return
puyuan1996
committed
b51efff
feature(pu): add textcraft env and configs
puyuan1996
committed
58a9824

Commits on Apr 28, 2026

add llm as policy and rlft ablations
xiongjyu
committed
68ae363
polish(pu): Align BabyAI config with ScalingInter-RL: multi-task 40 levels, per-level eval logging
puyuan1996
committed
1a59d6a
fix(pu): Fix BERT 1D tensor crash in eval, align BabyAI to 18 ScalingInter-RL levels

puyuan1996
and
claude
committed
f0c236b
add qwen as backbone
xiongjyu
committed
6793702
fix(pu): TP-pair drain in collector/eval, structured logging, eval-trajectory dump
puyuan1996
committed
96001b2
tmp
xiongjyu
committed
2309ab6

Commits on Apr 29, 2026

fix(pu): Fix CUDA OOM in chunked forward by removing hardcoded chunk_size floor, Fix dual CUDA OOM in train→vLLM handoff and large-batch forward pass
puyuan1996
committed
fe14fab

Commits on Apr 30, 2026

fix(pu): Fix LLM eval coverage, stabilize RFT training, add UniZero baseline config
puyuan1996
committed
fb53d67

Commits on May 2, 2026

fix(pu): Unify eval x-axis to env_step, add per-level UniZero evaluator, harden RFT training
puyuan1996
committed
41a212b

Commits on May 12, 2026

fix(pu): fix babyai_env.py
puyuan1996
committed
35d4c67

Commits on Jun 13, 2026

Commits on Jun 15, 2026