Hugging Face's logo

IffYuan
/
LeWorldModelplusplus

Robotics
JAX
English
world-model
planning
model-based-control
Model card Files
xet
 Community
IffYuan commited on
Commit
9fe532c
·
verified ·
1 Parent(s): 1ee9fcc

Add files using upload-large-folder tool

Browse files
Files changed (31) hide show
  1. lewm-control-suite/checkpoints/action-prior/cube/flags.json +66 -0
  2. lewm-control-suite/checkpoints/action-prior/cube/params_100000.pkl +3 -0
  3. lewm-control-suite/checkpoints/action-prior/pusht/flags.json +66 -0
  4. lewm-control-suite/checkpoints/action-prior/pusht/params_100000.pkl +3 -0
  5. lewm-control-suite/checkpoints/action-prior/reacher/flags.json +66 -0
  6. lewm-control-suite/checkpoints/action-prior/reacher/params_100000.pkl +3 -0
  7. lewm-control-suite/checkpoints/action-prior/tworoom/flags.json +66 -0
  8. lewm-control-suite/checkpoints/action-prior/tworoom/params_100000.pkl +3 -0
  9. lewm-control-suite/checkpoints/latent-path-flow/h25/cube/checkpoint_200000.msgpack +3 -0
  10. lewm-control-suite/checkpoints/latent-path-flow/h25/cube/config.json +46 -0
  11. lewm-control-suite/checkpoints/latent-path-flow/h25/pusht/checkpoint_200000.msgpack +3 -0
  12. lewm-control-suite/checkpoints/latent-path-flow/h25/pusht/config.json +46 -0
  13. lewm-control-suite/checkpoints/latent-path-flow/h25/reacher/checkpoint_200000.msgpack +3 -0
  14. lewm-control-suite/checkpoints/latent-path-flow/h25/reacher/config.json +46 -0
  15. lewm-control-suite/checkpoints/latent-path-flow/h25/tworoom/checkpoint_200000.msgpack +3 -0
  16. lewm-control-suite/checkpoints/latent-path-flow/h25/tworoom/config.json +46 -0
  17. lewm-control-suite/checkpoints/latent-path-flow/longh/cube/checkpoint_200000.msgpack +3 -0
  18. lewm-control-suite/checkpoints/latent-path-flow/longh/cube/config.json +45 -0
  19. lewm-control-suite/checkpoints/latent-path-flow/longh/pusht/checkpoint_200000.msgpack +3 -0
  20. lewm-control-suite/checkpoints/latent-path-flow/longh/pusht/config.json +45 -0
  21. lewm-control-suite/checkpoints/latent-path-flow/longh/reacher/checkpoint_200000.msgpack +3 -0
  22. lewm-control-suite/checkpoints/latent-path-flow/longh/reacher/config.json +45 -0
  23. lewm-control-suite/checkpoints/latent-path-flow/longh/tworoom/checkpoint_200000.msgpack +3 -0
  24. lewm-control-suite/checkpoints/latent-path-flow/longh/tworoom/config.json +45 -0
  25. lewm-control-suite/checkpoints/lewm/cube/config.json +34 -0
  26. lewm-control-suite/checkpoints/lewm/cube/weights_epoch_10.msgpack +3 -0
  27. lewm-control-suite/checkpoints/lewm/pusht/weights_epoch_10.msgpack +3 -0
  28. lewm-control-suite/checkpoints/lewm/reacher/config.json +34 -0
  29. lewm-control-suite/checkpoints/lewm/reacher/weights_epoch_10.msgpack +3 -0
  30. lewm-control-suite/checkpoints/lewm/tworoom/config.json +34 -0
  31. lewm-control-suite/checkpoints/lewm/tworoom/weights_epoch_10.msgpack +3 -0
lewm-control-suite/checkpoints/action-prior/cube/flags.json ADDED
@@ -0,0 +1,66 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "entrypoint": "train_gciql_chunk.py",
3
+ "env_name": null,
4
+ "dataset_path": "/data-training/yyf/datasets/lewm/cube_single_expert.lance",
5
+ "lewm_checkpoint": "/data-training/yyf/models/lewm-jax-seed3072/LeWMJAX_impala_lance_cube_single_bs128_e10_seed3072_fs5_h3_sigreg009_jpeg95/weights_epoch_10.msgpack",
6
+ "train_steps": 100000,
7
+ "seed": 777,
8
+ "agent": {
9
+ "actor_geom_sample": false,
10
+ "actor_hidden_dims": [
11
+ 512,
12
+ 512,
13
+ 512
14
+ ],
15
+ "actor_loss": "awr",
16
+ "actor_p_curgoal": 0.0,
17
+ "actor_p_randomgoal": 0.0,
18
+ "actor_p_trajgoal": 1.0,
19
+ "agent_name": "gciql_chunk_lewm",
20
+ "alpha": 3.0,
21
+ "batch_size": 256,
22
+ "chunk_size": 5,
23
+ "const_std": true,
24
+ "dataset_class": "GCChunkDataset",
25
+ "discount": 0.99,
26
+ "discrete": false,
27
+ "encoder": "impala_small",
28
+ "expectile": 0.9,
29
+ "frame_stack": null,
30
+ "gc_negative": true,
31
+ "latent_dim": 192,
32
+ "layer_norm": true,
33
+ "lr": 0.0003,
34
+ "p_aug": 0.0,
35
+ "representation_mode": "all",
36
+ "share_pi_encoder": true,
37
+ "share_q_encoder": true,
38
+ "share_v_encoder": true,
39
+ "tau": 0.005,
40
+ "value_geom_sample": true,
41
+ "value_hidden_dims": [
42
+ 512,
43
+ 512,
44
+ 512
45
+ ],
46
+ "value_p_curgoal": 0.2,
47
+ "value_p_randomgoal": 0.3,
48
+ "value_p_trajgoal": 0.5
49
+ },
50
+ "representation": {
51
+ "mode": "all",
52
+ "module": "LeWM.encode_pixels",
53
+ "output": "post_projector",
54
+ "latent_dim": 192,
55
+ "frozen": true,
56
+ "q": "lewm",
57
+ "v": "lewm",
58
+ "pi": "lewm",
59
+ "downstream_heads_shared": false,
60
+ "lewm_checkpoint": "/data-training/yyf/models/lewm-jax-seed3072/LeWMJAX_impala_lance_cube_single_bs128_e10_seed3072_fs5_h3_sigreg009_jpeg95/weights_epoch_10.msgpack"
61
+ },
62
+ "augmentation": {
63
+ "probability": 0.0,
64
+ "applied_before_all_encoders": true
65
+ }
66
+ }
lewm-control-suite/checkpoints/action-prior/cube/params_100000.pkl ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:90641c51023d509d17cf15a4a0a019ff32ea2572268fe0e41186fd939290a2b2
3
+ size 53007382
lewm-control-suite/checkpoints/action-prior/pusht/flags.json ADDED
@@ -0,0 +1,66 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "entrypoint": "train_gciql_chunk.py",
3
+ "env_name": null,
4
+ "dataset_path": "/data-training/yyf/datasets/lewm/pusht_expert_train.lance",
5
+ "lewm_checkpoint": "/data-training/yyf/models/lewm-jax-seed666/2026-08-19_23_LeWMJAX_impala_lance_pusht_expert_bs128_e10_seed666/weights_epoch_10.msgpack",
6
+ "train_steps": 100000,
7
+ "seed": 777,
8
+ "agent": {
9
+ "actor_geom_sample": false,
10
+ "actor_hidden_dims": [
11
+ 512,
12
+ 512,
13
+ 512
14
+ ],
15
+ "actor_loss": "awr",
16
+ "actor_p_curgoal": 0.0,
17
+ "actor_p_randomgoal": 0.0,
18
+ "actor_p_trajgoal": 1.0,
19
+ "agent_name": "gciql_chunk_lewm",
20
+ "alpha": 3.0,
21
+ "batch_size": 256,
22
+ "chunk_size": 5,
23
+ "const_std": true,
24
+ "dataset_class": "GCChunkDataset",
25
+ "discount": 0.99,
26
+ "discrete": false,
27
+ "encoder": "impala_small",
28
+ "expectile": 0.9,
29
+ "frame_stack": null,
30
+ "gc_negative": true,
31
+ "latent_dim": 192,
32
+ "layer_norm": true,
33
+ "lr": 0.0003,
34
+ "p_aug": 0.0,
35
+ "representation_mode": "all",
36
+ "share_pi_encoder": true,
37
+ "share_q_encoder": true,
38
+ "share_v_encoder": true,
39
+ "tau": 0.005,
40
+ "value_geom_sample": true,
41
+ "value_hidden_dims": [
42
+ 512,
43
+ 512,
44
+ 512
45
+ ],
46
+ "value_p_curgoal": 0.2,
47
+ "value_p_randomgoal": 0.3,
48
+ "value_p_trajgoal": 0.5
49
+ },
50
+ "representation": {
51
+ "mode": "all",
52
+ "module": "LeWM.encode_pixels",
53
+ "output": "post_projector",
54
+ "latent_dim": 192,
55
+ "frozen": true,
56
+ "q": "lewm",
57
+ "v": "lewm",
58
+ "pi": "lewm",
59
+ "downstream_heads_shared": false,
60
+ "lewm_checkpoint": "/data-training/yyf/models/lewm-jax-seed666/2026-08-19_23_LeWMJAX_impala_lance_pusht_expert_bs128_e10_seed666/weights_epoch_10.msgpack"
61
+ },
62
+ "augmentation": {
63
+ "probability": 0.0,
64
+ "applied_before_all_encoders": true
65
+ }
66
+ }
lewm-control-suite/checkpoints/action-prior/pusht/params_100000.pkl ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:3accd20f08fc014ba5e5c42fc9f6e086476b56b85692a8b25edfc948f977aa23
3
+ size 52546402
lewm-control-suite/checkpoints/action-prior/reacher/flags.json ADDED
@@ -0,0 +1,66 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "entrypoint": "train_gciql_chunk.py",
3
+ "env_name": null,
4
+ "dataset_path": "/data-training/yyf/datasets/lewm/reacher.lance",
5
+ "lewm_checkpoint": "/data-training/yyf/models/lewm-jax-seed3072/LeWMJAX_impala_lance_reacher_bs128_e10_seed3072_fs5_h3_sigreg009_jpeg95/weights_epoch_10.msgpack",
6
+ "train_steps": 100000,
7
+ "seed": 777,
8
+ "agent": {
9
+ "actor_geom_sample": false,
10
+ "actor_hidden_dims": [
11
+ 512,
12
+ 512,
13
+ 512
14
+ ],
15
+ "actor_loss": "awr",
16
+ "actor_p_curgoal": 0.0,
17
+ "actor_p_randomgoal": 0.0,
18
+ "actor_p_trajgoal": 1.0,
19
+ "agent_name": "gciql_chunk_lewm",
20
+ "alpha": 3.0,
21
+ "batch_size": 256,
22
+ "chunk_size": 5,
23
+ "const_std": true,
24
+ "dataset_class": "GCChunkDataset",
25
+ "discount": 0.99,
26
+ "discrete": false,
27
+ "encoder": "impala_small",
28
+ "expectile": 0.9,
29
+ "frame_stack": null,
30
+ "gc_negative": true,
31
+ "latent_dim": 192,
32
+ "layer_norm": true,
33
+ "lr": 0.0003,
34
+ "p_aug": 0.0,
35
+ "representation_mode": "all",
36
+ "share_pi_encoder": true,
37
+ "share_q_encoder": true,
38
+ "share_v_encoder": true,
39
+ "tau": 0.005,
40
+ "value_geom_sample": true,
41
+ "value_hidden_dims": [
42
+ 512,
43
+ 512,
44
+ 512
45
+ ],
46
+ "value_p_curgoal": 0.2,
47
+ "value_p_randomgoal": 0.3,
48
+ "value_p_trajgoal": 0.5
49
+ },
50
+ "representation": {
51
+ "mode": "all",
52
+ "module": "LeWM.encode_pixels",
53
+ "output": "post_projector",
54
+ "latent_dim": 192,
55
+ "frozen": true,
56
+ "q": "lewm",
57
+ "v": "lewm",
58
+ "pi": "lewm",
59
+ "downstream_heads_shared": false,
60
+ "lewm_checkpoint": "/data-training/yyf/models/lewm-jax-seed3072/LeWMJAX_impala_lance_reacher_bs128_e10_seed3072_fs5_h3_sigreg009_jpeg95/weights_epoch_10.msgpack"
61
+ },
62
+ "augmentation": {
63
+ "probability": 0.0,
64
+ "applied_before_all_encoders": true
65
+ }
66
+ }
lewm-control-suite/checkpoints/action-prior/reacher/params_100000.pkl ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:9693ceb0f078d114227ff793f8863e8d22fbdfdb0410f46732179e1dcb2a4d0d
3
+ size 52546402
lewm-control-suite/checkpoints/action-prior/tworoom/flags.json ADDED
@@ -0,0 +1,66 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "entrypoint": "train_gciql_chunk.py",
3
+ "env_name": null,
4
+ "dataset_path": "/data-training/yyf/datasets/lewm/tworoom.lance",
5
+ "lewm_checkpoint": "/data-training/yyf/models/lewm-jax-seed3072/LeWMJAX_impala_lance_tworoom_bs128_e10_seed3072_fs5_h3_sigreg009_jpeg95/weights_epoch_10.msgpack",
6
+ "train_steps": 100000,
7
+ "seed": 777,
8
+ "agent": {
9
+ "actor_geom_sample": false,
10
+ "actor_hidden_dims": [
11
+ 512,
12
+ 512,
13
+ 512
14
+ ],
15
+ "actor_loss": "awr",
16
+ "actor_p_curgoal": 0.0,
17
+ "actor_p_randomgoal": 0.0,
18
+ "actor_p_trajgoal": 1.0,
19
+ "agent_name": "gciql_chunk_lewm",
20
+ "alpha": 3.0,
21
+ "batch_size": 256,
22
+ "chunk_size": 5,
23
+ "const_std": true,
24
+ "dataset_class": "GCChunkDataset",
25
+ "discount": 0.99,
26
+ "discrete": false,
27
+ "encoder": "impala_small",
28
+ "expectile": 0.9,
29
+ "frame_stack": null,
30
+ "gc_negative": true,
31
+ "latent_dim": 192,
32
+ "layer_norm": true,
33
+ "lr": 0.0003,
34
+ "p_aug": 0.0,
35
+ "representation_mode": "all",
36
+ "share_pi_encoder": true,
37
+ "share_q_encoder": true,
38
+ "share_v_encoder": true,
39
+ "tau": 0.005,
40
+ "value_geom_sample": true,
41
+ "value_hidden_dims": [
42
+ 512,
43
+ 512,
44
+ 512
45
+ ],
46
+ "value_p_curgoal": 0.2,
47
+ "value_p_randomgoal": 0.3,
48
+ "value_p_trajgoal": 0.5
49
+ },
50
+ "representation": {
51
+ "mode": "all",
52
+ "module": "LeWM.encode_pixels",
53
+ "output": "post_projector",
54
+ "latent_dim": 192,
55
+ "frozen": true,
56
+ "q": "lewm",
57
+ "v": "lewm",
58
+ "pi": "lewm",
59
+ "downstream_heads_shared": false,
60
+ "lewm_checkpoint": "/data-training/yyf/models/lewm-jax-seed3072/LeWMJAX_impala_lance_tworoom_bs128_e10_seed3072_fs5_h3_sigreg009_jpeg95/weights_epoch_10.msgpack"
61
+ },
62
+ "augmentation": {
63
+ "probability": 0.0,
64
+ "applied_before_all_encoders": true
65
+ }
66
+ }
lewm-control-suite/checkpoints/action-prior/tworoom/params_100000.pkl ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:7ac214436cefbe646e0fa39a40dc538ba83db042dd96a77b18a40ad5de8e9ab0
3
+ size 52546402
lewm-control-suite/checkpoints/latent-path-flow/h25/cube/checkpoint_200000.msgpack ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:1aa7e8671525dfb553fa94236ad5b8ab1f628807b2a0d1592890c9b0e9b81b74
3
+ size 299898830
lewm-control-suite/checkpoints/latent-path-flow/h25/cube/config.json ADDED
@@ -0,0 +1,46 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "action_block": 5,
3
+ "architecture": "latent_path_flow_transformer_encoder",
4
+ "batch_size": 1024,
5
+ "checkpoint_interval": 25000,
6
+ "conditioning": "history_goal_time_adaln",
7
+ "depth": 4,
8
+ "ema_decay": 0.9999,
9
+ "embed_dim": 192,
10
+ "eval_batch_size": 1024,
11
+ "eval_interval": 10000,
12
+ "exp_name": "latent_pathflow_cube_lewm3072_hist3_sg10_ab5_goalstride5_goalmax25_cfm_ns8_n200000_b1024_s0",
13
+ "ff_dim": 2048,
14
+ "final_learning_rate": 1e-05,
15
+ "flow_sampling_steps": 16,
16
+ "flow_solver": "euler",
17
+ "goal_sampling": "uniform_distance_first_aligned_future_same_trajectory_stride_5_max_25",
18
+ "gradient_clip": 1.0,
19
+ "hidden_dim": 512,
20
+ "history_size": 3,
21
+ "latent_dataset": "/data-training/yyf/datasets/lewm-latents/cube_single_expert__lewm_s3072_e10_z192.h5",
22
+ "learning_rate": 0.0001,
23
+ "lewm_checkpoint_sha256": "0707653d129c2766aa161ead7f777852a783fc82bd940706623b14b80e6d277d",
24
+ "log_interval": 1000,
25
+ "loss": "conditional_path_flow_matching_mse",
26
+ "max_goal_steps": 25,
27
+ "num_episodes": 10000,
28
+ "num_heads": 8,
29
+ "num_rows": 2010000,
30
+ "num_samples": 8,
31
+ "num_train_episodes": 9500,
32
+ "num_train_transitions": 1862000,
33
+ "num_val_episodes": 500,
34
+ "num_val_transitions": 98000,
35
+ "resume": true,
36
+ "save_dir": "/data-training/yyf/ogbench-lewm-policy-runs/latent-path-flow-k10-goalmax25/latent_pathflow_cube_lewm3072_hist3_sg10_ab5_goalstride5_goalmax25_cfm_ns8_n200000_b1024_s0",
37
+ "seed": 0,
38
+ "split_seed": 0,
39
+ "subgoal_steps": 10,
40
+ "time_dim": 64,
41
+ "train_fraction": 0.95,
42
+ "train_steps": 200000,
43
+ "validation_pairs": 10000,
44
+ "warmup_steps": 5000,
45
+ "weight_decay": 0.0001
46
+ }
lewm-control-suite/checkpoints/latent-path-flow/h25/pusht/checkpoint_200000.msgpack ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:2ce4bbd4684c668301a296c021afd7aefb3bb4da08e1f39808c439fe2d8b1dff
3
+ size 299898830
lewm-control-suite/checkpoints/latent-path-flow/h25/pusht/config.json ADDED
@@ -0,0 +1,46 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "action_block": 5,
3
+ "architecture": "latent_path_flow_transformer_encoder",
4
+ "batch_size": 1024,
5
+ "checkpoint_interval": 25000,
6
+ "conditioning": "history_goal_time_adaln",
7
+ "depth": 4,
8
+ "ema_decay": 0.9999,
9
+ "embed_dim": 192,
10
+ "eval_batch_size": 1024,
11
+ "eval_interval": 10000,
12
+ "exp_name": "latent_pathflow_pusht_lewm666_hist3_sg10_ab5_goalstride5_goalmax25_cfm_ns8_n200000_b1024_s0",
13
+ "ff_dim": 2048,
14
+ "final_learning_rate": 1e-05,
15
+ "flow_sampling_steps": 16,
16
+ "flow_solver": "euler",
17
+ "goal_sampling": "uniform_distance_first_aligned_future_same_trajectory_stride_5_max_25",
18
+ "gradient_clip": 1.0,
19
+ "hidden_dim": 512,
20
+ "history_size": 3,
21
+ "latent_dataset": "/data-training/yyf/datasets/lewm-latents/pusht_expert_train__lewm_s666_e10_z192.h5",
22
+ "learning_rate": 0.0001,
23
+ "lewm_checkpoint_sha256": "ee02a100decff720e0fbd757a60e9281db0159ec04d784735fb22cfe568f3090",
24
+ "log_interval": 1000,
25
+ "loss": "conditional_path_flow_matching_mse",
26
+ "max_goal_steps": 25,
27
+ "num_episodes": 18685,
28
+ "num_heads": 8,
29
+ "num_rows": 2336736,
30
+ "num_samples": 8,
31
+ "num_train_episodes": 17750,
32
+ "num_train_transitions": 2130385,
33
+ "num_val_episodes": 935,
34
+ "num_val_transitions": 112926,
35
+ "resume": true,
36
+ "save_dir": "/data-training/yyf/ogbench-lewm-policy-runs/latent-path-flow-k10-goalmax25/latent_pathflow_pusht_lewm666_hist3_sg10_ab5_goalstride5_goalmax25_cfm_ns8_n200000_b1024_s0",
37
+ "seed": 0,
38
+ "split_seed": 0,
39
+ "subgoal_steps": 10,
40
+ "time_dim": 64,
41
+ "train_fraction": 0.95,
42
+ "train_steps": 200000,
43
+ "validation_pairs": 10000,
44
+ "warmup_steps": 5000,
45
+ "weight_decay": 0.0001
46
+ }
lewm-control-suite/checkpoints/latent-path-flow/h25/reacher/checkpoint_200000.msgpack ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:b75bb1d9d7c4b5418994a57048efa63cfa428934a373e6355763e42306644d51
3
+ size 299898830
lewm-control-suite/checkpoints/latent-path-flow/h25/reacher/config.json ADDED
@@ -0,0 +1,46 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "action_block": 5,
3
+ "architecture": "latent_path_flow_transformer_encoder",
4
+ "batch_size": 1024,
5
+ "checkpoint_interval": 25000,
6
+ "conditioning": "history_goal_time_adaln",
7
+ "depth": 4,
8
+ "ema_decay": 0.9999,
9
+ "embed_dim": 192,
10
+ "eval_batch_size": 1024,
11
+ "eval_interval": 10000,
12
+ "exp_name": "latent_pathflow_reacher_lewm3072_hist3_sg10_ab5_goalstride5_goalmax25_cfm_ns8_n200000_b1024_s0",
13
+ "ff_dim": 2048,
14
+ "final_learning_rate": 1e-05,
15
+ "flow_sampling_steps": 16,
16
+ "flow_solver": "euler",
17
+ "goal_sampling": "uniform_distance_first_aligned_future_same_trajectory_stride_5_max_25",
18
+ "gradient_clip": 1.0,
19
+ "hidden_dim": 512,
20
+ "history_size": 3,
21
+ "latent_dataset": "/data-training/yyf/datasets/lewm-latents/reacher__lewm_s3072_e10_z192.h5",
22
+ "learning_rate": 0.0001,
23
+ "lewm_checkpoint_sha256": "91f888e5de8f7872b18f209ddbaff0f0f03644d2ef66429c0b09c0d6c34cee3a",
24
+ "log_interval": 1000,
25
+ "loss": "conditional_path_flow_matching_mse",
26
+ "max_goal_steps": 25,
27
+ "num_episodes": 10000,
28
+ "num_heads": 8,
29
+ "num_rows": 2010000,
30
+ "num_samples": 8,
31
+ "num_train_episodes": 9500,
32
+ "num_train_transitions": 1862000,
33
+ "num_val_episodes": 500,
34
+ "num_val_transitions": 98000,
35
+ "resume": true,
36
+ "save_dir": "/data-training/yyf/ogbench-lewm-policy-runs/latent-path-flow-k10-goalmax25/latent_pathflow_reacher_lewm3072_hist3_sg10_ab5_goalstride5_goalmax25_cfm_ns8_n200000_b1024_s0",
37
+ "seed": 0,
38
+ "split_seed": 0,
39
+ "subgoal_steps": 10,
40
+ "time_dim": 64,
41
+ "train_fraction": 0.95,
42
+ "train_steps": 200000,
43
+ "validation_pairs": 10000,
44
+ "warmup_steps": 5000,
45
+ "weight_decay": 0.0001
46
+ }
lewm-control-suite/checkpoints/latent-path-flow/h25/tworoom/checkpoint_200000.msgpack ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:d79adde1b5d490b37810baaf842c95411d9de013e9c813b570c64746e4fdaf9e
3
+ size 299898830
lewm-control-suite/checkpoints/latent-path-flow/h25/tworoom/config.json ADDED
@@ -0,0 +1,46 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "action_block": 5,
3
+ "architecture": "latent_path_flow_transformer_encoder",
4
+ "batch_size": 1024,
5
+ "checkpoint_interval": 25000,
6
+ "conditioning": "history_goal_time_adaln",
7
+ "depth": 4,
8
+ "ema_decay": 0.9999,
9
+ "embed_dim": 192,
10
+ "eval_batch_size": 1024,
11
+ "eval_interval": 10000,
12
+ "exp_name": "latent_pathflow_tworoom_lewm3072_hist3_sg10_ab5_goalstride5_goalmax25_cfm_ns8_n200000_b1024_s0",
13
+ "ff_dim": 2048,
14
+ "final_learning_rate": 1e-05,
15
+ "flow_sampling_steps": 16,
16
+ "flow_solver": "euler",
17
+ "goal_sampling": "uniform_distance_first_aligned_future_same_trajectory_stride_5_max_25",
18
+ "gradient_clip": 1.0,
19
+ "hidden_dim": 512,
20
+ "history_size": 3,
21
+ "latent_dataset": "/data-training/yyf/datasets/lewm-latents/tworoom__lewm_s3072_e10_z192.h5",
22
+ "learning_rate": 0.0001,
23
+ "lewm_checkpoint_sha256": "e77313033e83b41be36ae0c18880b758c4169c102af5bd357fe041b181836e46",
24
+ "log_interval": 1000,
25
+ "loss": "conditional_path_flow_matching_mse",
26
+ "max_goal_steps": 25,
27
+ "num_episodes": 10000,
28
+ "num_heads": 8,
29
+ "num_rows": 920809,
30
+ "num_samples": 8,
31
+ "num_train_episodes": 9500,
32
+ "num_train_transitions": 827091,
33
+ "num_val_episodes": 500,
34
+ "num_val_transitions": 43718,
35
+ "resume": true,
36
+ "save_dir": "/data-training/yyf/ogbench-lewm-policy-runs/latent-path-flow-k10-goalmax25/latent_pathflow_tworoom_lewm3072_hist3_sg10_ab5_goalstride5_goalmax25_cfm_ns8_n200000_b1024_s0",
37
+ "seed": 0,
38
+ "split_seed": 0,
39
+ "subgoal_steps": 10,
40
+ "time_dim": 64,
41
+ "train_fraction": 0.95,
42
+ "train_steps": 200000,
43
+ "validation_pairs": 10000,
44
+ "warmup_steps": 5000,
45
+ "weight_decay": 0.0001
46
+ }
lewm-control-suite/checkpoints/latent-path-flow/longh/cube/checkpoint_200000.msgpack ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:5fe72a2c09fc11cdb4966f57f7b18aea0501e130cad7357fe32a6d8bb04c30f6
3
+ size 299898830
lewm-control-suite/checkpoints/latent-path-flow/longh/cube/config.json ADDED
@@ -0,0 +1,45 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "action_block": 5,
3
+ "architecture": "latent_path_flow_transformer_encoder",
4
+ "batch_size": 1024,
5
+ "checkpoint_interval": 25000,
6
+ "conditioning": "history_goal_time_adaln",
7
+ "depth": 4,
8
+ "ema_decay": 0.9999,
9
+ "embed_dim": 192,
10
+ "eval_batch_size": 1024,
11
+ "eval_interval": 10000,
12
+ "exp_name": "latent_pathflow_cube_lewm3072_hist3_sg10_ab5_cfm_ns8_n200000_b1024_s0",
13
+ "ff_dim": 2048,
14
+ "final_learning_rate": 1e-05,
15
+ "flow_sampling_steps": 16,
16
+ "flow_solver": "euler",
17
+ "goal_sampling": "hiql_uniform_future_same_trajectory",
18
+ "gradient_clip": 1.0,
19
+ "hidden_dim": 512,
20
+ "history_size": 3,
21
+ "latent_dataset": "/data-training/yyf/datasets/lewm-latents/cube_single_expert__lewm_s3072_e10_z192.h5",
22
+ "learning_rate": 0.0001,
23
+ "lewm_checkpoint_sha256": "0707653d129c2766aa161ead7f777852a783fc82bd940706623b14b80e6d277d",
24
+ "log_interval": 1000,
25
+ "loss": "conditional_path_flow_matching_mse",
26
+ "num_episodes": 10000,
27
+ "num_heads": 8,
28
+ "num_rows": 2010000,
29
+ "num_samples": 8,
30
+ "num_train_episodes": 9500,
31
+ "num_train_transitions": 1900000,
32
+ "num_val_episodes": 500,
33
+ "num_val_transitions": 100000,
34
+ "resume": true,
35
+ "save_dir": "/data-training/yyf/ogbench-lewm-policy-runs/latent-path-flow-k10/latent_pathflow_cube_lewm3072_hist3_sg10_ab5_cfm_ns8_n200000_b1024_s0",
36
+ "seed": 0,
37
+ "split_seed": 0,
38
+ "subgoal_steps": 10,
39
+ "time_dim": 64,
40
+ "train_fraction": 0.95,
41
+ "train_steps": 200000,
42
+ "validation_pairs": 10000,
43
+ "warmup_steps": 5000,
44
+ "weight_decay": 0.0001
45
+ }
lewm-control-suite/checkpoints/latent-path-flow/longh/pusht/checkpoint_200000.msgpack ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:d130840ede60a8739d71d47b05c1c41996ea1b2053df358d8f27378349135a74
3
+ size 299898830
lewm-control-suite/checkpoints/latent-path-flow/longh/pusht/config.json ADDED
@@ -0,0 +1,45 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "action_block": 5,
3
+ "architecture": "latent_path_flow_transformer_encoder",
4
+ "batch_size": 1024,
5
+ "checkpoint_interval": 25000,
6
+ "conditioning": "history_goal_time_adaln",
7
+ "depth": 4,
8
+ "ema_decay": 0.9999,
9
+ "embed_dim": 192,
10
+ "eval_batch_size": 1024,
11
+ "eval_interval": 10000,
12
+ "exp_name": "latent_pathflow_pusht_lewm666_hist3_sg10_ab5_cfm_ns8_n200000_b1024_s0",
13
+ "ff_dim": 2048,
14
+ "final_learning_rate": 1e-05,
15
+ "flow_sampling_steps": 16,
16
+ "flow_solver": "euler",
17
+ "goal_sampling": "hiql_uniform_future_same_trajectory",
18
+ "gradient_clip": 1.0,
19
+ "hidden_dim": 512,
20
+ "history_size": 3,
21
+ "latent_dataset": "/data-training/yyf/datasets/lewm-latents/pusht_expert_train__lewm_s666_e10_z192.h5",
22
+ "learning_rate": 0.0001,
23
+ "lewm_checkpoint_sha256": "ee02a100decff720e0fbd757a60e9281db0159ec04d784735fb22cfe568f3090",
24
+ "log_interval": 1000,
25
+ "loss": "conditional_path_flow_matching_mse",
26
+ "num_episodes": 18685,
27
+ "num_heads": 8,
28
+ "num_rows": 2336736,
29
+ "num_samples": 8,
30
+ "num_train_episodes": 17750,
31
+ "num_train_transitions": 2201385,
32
+ "num_val_episodes": 935,
33
+ "num_val_transitions": 116666,
34
+ "resume": true,
35
+ "save_dir": "/data-training/yyf/ogbench-lewm-policy-runs/latent-path-flow-k10/latent_pathflow_pusht_lewm666_hist3_sg10_ab5_cfm_ns8_n200000_b1024_s0",
36
+ "seed": 0,
37
+ "split_seed": 0,
38
+ "subgoal_steps": 10,
39
+ "time_dim": 64,
40
+ "train_fraction": 0.95,
41
+ "train_steps": 200000,
42
+ "validation_pairs": 10000,
43
+ "warmup_steps": 5000,
44
+ "weight_decay": 0.0001
45
+ }
lewm-control-suite/checkpoints/latent-path-flow/longh/reacher/checkpoint_200000.msgpack ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ac7b00dec99270bed8659399484ce46726027eb3cb5eea2e2388a44d2e824144
3
+ size 299898830
lewm-control-suite/checkpoints/latent-path-flow/longh/reacher/config.json ADDED
@@ -0,0 +1,45 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "action_block": 5,
3
+ "architecture": "latent_path_flow_transformer_encoder",
4
+ "batch_size": 1024,
5
+ "checkpoint_interval": 25000,
6
+ "conditioning": "history_goal_time_adaln",
7
+ "depth": 4,
8
+ "ema_decay": 0.9999,
9
+ "embed_dim": 192,
10
+ "eval_batch_size": 1024,
11
+ "eval_interval": 10000,
12
+ "exp_name": "latent_pathflow_reacher_lewm3072_hist3_sg10_ab5_cfm_ns8_n200000_b1024_s0",
13
+ "ff_dim": 2048,
14
+ "final_learning_rate": 1e-05,
15
+ "flow_sampling_steps": 16,
16
+ "flow_solver": "euler",
17
+ "goal_sampling": "hiql_uniform_future_same_trajectory",
18
+ "gradient_clip": 1.0,
19
+ "hidden_dim": 512,
20
+ "history_size": 3,
21
+ "latent_dataset": "/data-training/yyf/datasets/lewm-latents/reacher__lewm_s3072_e10_z192.h5",
22
+ "learning_rate": 0.0001,
23
+ "lewm_checkpoint_sha256": "91f888e5de8f7872b18f209ddbaff0f0f03644d2ef66429c0b09c0d6c34cee3a",
24
+ "log_interval": 1000,
25
+ "loss": "conditional_path_flow_matching_mse",
26
+ "num_episodes": 10000,
27
+ "num_heads": 8,
28
+ "num_rows": 2010000,
29
+ "num_samples": 8,
30
+ "num_train_episodes": 9500,
31
+ "num_train_transitions": 1900000,
32
+ "num_val_episodes": 500,
33
+ "num_val_transitions": 100000,
34
+ "resume": true,
35
+ "save_dir": "/data-training/yyf/ogbench-lewm-policy-runs/latent-path-flow-k10/latent_pathflow_reacher_lewm3072_hist3_sg10_ab5_cfm_ns8_n200000_b1024_s0",
36
+ "seed": 0,
37
+ "split_seed": 0,
38
+ "subgoal_steps": 10,
39
+ "time_dim": 64,
40
+ "train_fraction": 0.95,
41
+ "train_steps": 200000,
42
+ "validation_pairs": 10000,
43
+ "warmup_steps": 5000,
44
+ "weight_decay": 0.0001
45
+ }
lewm-control-suite/checkpoints/latent-path-flow/longh/tworoom/checkpoint_200000.msgpack ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:644b2b6fbf775983975c621cb2115d8c82e54c820d68cc213f16f44ed206a3d8
3
+ size 299898830
lewm-control-suite/checkpoints/latent-path-flow/longh/tworoom/config.json ADDED
@@ -0,0 +1,45 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "action_block": 5,
3
+ "architecture": "latent_path_flow_transformer_encoder",
4
+ "batch_size": 1024,
5
+ "checkpoint_interval": 25000,
6
+ "conditioning": "history_goal_time_adaln",
7
+ "depth": 4,
8
+ "ema_decay": 0.9999,
9
+ "embed_dim": 192,
10
+ "eval_batch_size": 1024,
11
+ "eval_interval": 10000,
12
+ "exp_name": "latent_pathflow_tworoom_lewm3072_hist3_sg10_ab5_cfm_ns8_n200000_b1024_s0",
13
+ "ff_dim": 2048,
14
+ "final_learning_rate": 1e-05,
15
+ "flow_sampling_steps": 16,
16
+ "flow_solver": "euler",
17
+ "goal_sampling": "hiql_uniform_future_same_trajectory",
18
+ "gradient_clip": 1.0,
19
+ "hidden_dim": 512,
20
+ "history_size": 3,
21
+ "latent_dataset": "/data-training/yyf/datasets/lewm-latents/tworoom__lewm_s3072_e10_z192.h5",
22
+ "learning_rate": 0.0001,
23
+ "lewm_checkpoint_sha256": "e77313033e83b41be36ae0c18880b758c4169c102af5bd357fe041b181836e46",
24
+ "log_interval": 1000,
25
+ "loss": "conditional_path_flow_matching_mse",
26
+ "num_episodes": 10000,
27
+ "num_heads": 8,
28
+ "num_rows": 920809,
29
+ "num_samples": 8,
30
+ "num_train_episodes": 9500,
31
+ "num_train_transitions": 865091,
32
+ "num_val_episodes": 500,
33
+ "num_val_transitions": 45718,
34
+ "resume": true,
35
+ "save_dir": "/data-training/yyf/ogbench-lewm-policy-runs/latent-path-flow-k10/latent_pathflow_tworoom_lewm3072_hist3_sg10_ab5_cfm_ns8_n200000_b1024_s0",
36
+ "seed": 0,
37
+ "split_seed": 0,
38
+ "subgoal_steps": 10,
39
+ "time_dim": 64,
40
+ "train_fraction": 0.95,
41
+ "train_steps": 200000,
42
+ "validation_pairs": 10000,
43
+ "warmup_steps": 5000,
44
+ "weight_decay": 0.0001
45
+ }
lewm-control-suite/checkpoints/lewm/cube/config.json ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "exp_name": "LeWMJAX_impala_lance_cube_single_bs128_e10_seed3072_fs5_h3_sigreg009_jpeg95",
3
+ "dataset_path": "/data/dzb/stablewm-data/datasets/cube_single_expert.lance",
4
+ "architecture": "lewm_impala_small",
5
+ "encoder": "impala_small",
6
+ "seed": 3072,
7
+ "epochs": 10,
8
+ "batch_size": 128,
9
+ "decode_workers": 6,
10
+ "train_fraction": 0.9,
11
+ "image_size": 224,
12
+ "embed_dim": 192,
13
+ "patch_size": 14,
14
+ "history_size": 3,
15
+ "num_preds": 1,
16
+ "frameskip": 5,
17
+ "projector_hidden_dim": 2048,
18
+ "action_smoothed_dim": 10,
19
+ "action_mlp_scale": 4,
20
+ "predictor_depth": 6,
21
+ "predictor_heads": 16,
22
+ "predictor_dim_head": 64,
23
+ "predictor_mlp_dim": 2048,
24
+ "predictor_dropout": 0.1,
25
+ "predictor_emb_dropout": 0.0,
26
+ "learning_rate": 5e-05,
27
+ "lr_schedule": "optimizer_step_warmup_cosine_1pct",
28
+ "weight_decay": 0.001,
29
+ "gradient_clip": 1.0,
30
+ "sigreg_weight": 0.09,
31
+ "sigreg_knots": 17,
32
+ "sigreg_num_proj": 1024,
33
+ "precision": "bf16"
34
+ }
lewm-control-suite/checkpoints/lewm/cube/weights_epoch_10.msgpack ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:0707653d129c2766aa161ead7f777852a783fc82bd940706623b14b80e6d277d
3
+ size 104398277
lewm-control-suite/checkpoints/lewm/pusht/weights_epoch_10.msgpack ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ee02a100decff720e0fbd757a60e9281db0159ec04d784735fb22cfe568f3090
3
+ size 104397665
lewm-control-suite/checkpoints/lewm/reacher/config.json ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "exp_name": "LeWMJAX_impala_lance_reacher_bs128_e10_seed3072_fs5_h3_sigreg009_jpeg95",
3
+ "dataset_path": "/data/dzb/stablewm-data/datasets/reacher.lance",
4
+ "architecture": "lewm_impala_small",
5
+ "encoder": "impala_small",
6
+ "seed": 3072,
7
+ "epochs": 10,
8
+ "batch_size": 128,
9
+ "decode_workers": 6,
10
+ "train_fraction": 0.9,
11
+ "image_size": 224,
12
+ "embed_dim": 192,
13
+ "patch_size": 14,
14
+ "history_size": 3,
15
+ "num_preds": 1,
16
+ "frameskip": 5,
17
+ "projector_hidden_dim": 2048,
18
+ "action_smoothed_dim": 10,
19
+ "action_mlp_scale": 4,
20
+ "predictor_depth": 6,
21
+ "predictor_heads": 16,
22
+ "predictor_dim_head": 64,
23
+ "predictor_mlp_dim": 2048,
24
+ "predictor_dropout": 0.1,
25
+ "predictor_emb_dropout": 0.0,
26
+ "learning_rate": 5e-05,
27
+ "lr_schedule": "optimizer_step_warmup_cosine_1pct",
28
+ "weight_decay": 0.001,
29
+ "gradient_clip": 1.0,
30
+ "sigreg_weight": 0.09,
31
+ "sigreg_knots": 17,
32
+ "sigreg_num_proj": 1024,
33
+ "precision": "bf16"
34
+ }
lewm-control-suite/checkpoints/lewm/reacher/weights_epoch_10.msgpack ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:91f888e5de8f7872b18f209ddbaff0f0f03644d2ef66429c0b09c0d6c34cee3a
3
+ size 104397677
lewm-control-suite/checkpoints/lewm/tworoom/config.json ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "exp_name": "LeWMJAX_impala_lance_tworoom_bs128_e10_seed3072_fs5_h3_sigreg009_jpeg95",
3
+ "dataset_path": "/data/dzb/stablewm-data/datasets/tworoom.lance",
4
+ "architecture": "lewm_impala_small",
5
+ "encoder": "impala_small",
6
+ "seed": 3072,
7
+ "epochs": 10,
8
+ "batch_size": 128,
9
+ "decode_workers": 6,
10
+ "train_fraction": 0.9,
11
+ "image_size": 224,
12
+ "embed_dim": 192,
13
+ "patch_size": 14,
14
+ "history_size": 3,
15
+ "num_preds": 1,
16
+ "frameskip": 5,
17
+ "projector_hidden_dim": 2048,
18
+ "action_smoothed_dim": 10,
19
+ "action_mlp_scale": 4,
20
+ "predictor_depth": 6,
21
+ "predictor_heads": 16,
22
+ "predictor_dim_head": 64,
23
+ "predictor_mlp_dim": 2048,
24
+ "predictor_dropout": 0.1,
25
+ "predictor_emb_dropout": 0.0,
26
+ "learning_rate": 5e-05,
27
+ "lr_schedule": "optimizer_step_warmup_cosine_1pct",
28
+ "weight_decay": 0.001,
29
+ "gradient_clip": 1.0,
30
+ "sigreg_weight": 0.09,
31
+ "sigreg_knots": 17,
32
+ "sigreg_num_proj": 1024,
33
+ "precision": "bf16"
34
+ }
lewm-control-suite/checkpoints/lewm/tworoom/weights_epoch_10.msgpack ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:e77313033e83b41be36ae0c18880b758c4169c102af5bd357fe041b181836e46
3
+ size 104397677