Upload checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_mse_only/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_mse_only

checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_mse_only/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_mse_only/wandb/debug-internal.log

@@ -1,13 +1,8 @@
-{"time":"2026-01-02T21:43:05.63786614Z","level":"INFO","msg":"stream: starting","core version":"0.23.1"}
-{"time":"2026-01-02T21:43:06.061238769Z","level":"WARN","msg":"featurechecker: GraphQL client is nil, skipping feature loading"}
-{"time":"2026-01-02T21:43:06.061320164Z","level":"INFO","msg":"stream: created new stream","id":"vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0"}
-{"time":"2026-01-02T21:43:06.061627231Z","level":"INFO","msg":"handler: started","stream_id":"vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0"}
-{"time":"2026-01-02T21:43:06.061722049Z","level":"INFO","msg":"stream: started","id":"vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0"}
-{"time":"2026-01-02T21:43:06.061752682Z","level":"INFO","msg":"writer: started","stream_id":"vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0"}
-{"time":"2026-01-02T21:43:06.061765788Z","level":"INFO","msg":"sender: started","stream_id":"vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0"}
-{"time":"2026-01-02T21:43:06.062871122Z","level":"WARN","msg":"runupserter: server does not expand metric globs but the x_server_side_expand_glob_metrics setting is set; ignoring"}
-{"time":"2026-01-03T17:12:43.962447603Z","level":"INFO","msg":"handler: operation stats","stats":{}}
-{"time":"2026-01-03T17:12:43.96605726Z","level":"INFO","msg":"stream: closing","id":"vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0"}
-{"time":"2026-01-03T17:12:43.966072728Z","level":"INFO","msg":"handler: closed","stream_id":"vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0"}
-{"time":"2026-01-03T17:12:43.966105326Z","level":"INFO","msg":"sender: closed","stream_id":"vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0"}
-{"time":"2026-01-03T17:12:43.966121948Z","level":"INFO","msg":"stream: closed","id":"vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0"}
+{"time":"2026-01-04T09:32:01.115692117Z","level":"INFO","msg":"stream: starting","core version":"0.23.1"}
+{"time":"2026-01-04T09:32:01.44318528Z","level":"WARN","msg":"featurechecker: GraphQL client is nil, skipping feature loading"}
+{"time":"2026-01-04T09:32:01.443259623Z","level":"INFO","msg":"stream: created new stream","id":"vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0"}
+{"time":"2026-01-04T09:32:01.443726951Z","level":"INFO","msg":"stream: started","id":"vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0"}
+{"time":"2026-01-04T09:32:01.44372382Z","level":"INFO","msg":"handler: started","stream_id":"vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0"}
+{"time":"2026-01-04T09:32:01.443767451Z","level":"INFO","msg":"writer: started","stream_id":"vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0"}
+{"time":"2026-01-04T09:32:01.443797476Z","level":"INFO","msg":"sender: started","stream_id":"vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0"}
+{"time":"2026-01-04T09:32:01.444535545Z","level":"WARN","msg":"runupserter: server does not expand metric globs but the x_server_side_expand_glob_metrics setting is set; ignoring"}

checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_mse_only/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_mse_only/wandb/debug.log

@@ -1,28 +1,24 @@
-2026-01-02 21:43:04,973 INFO    MainThread:13227 [wandb_setup.py:_flush():80] Current SDK version is 0.23.1
-2026-01-02 21:43:04,973 INFO    MainThread:13227 [wandb_setup.py:_flush():80] Configure stats pid to 13227
-2026-01-02 21:43:04,973 INFO    MainThread:13227 [wandb_setup.py:_flush():80] Loading settings from /home/clouduser/.config/wandb/settings
-2026-01-02 21:43:04,973 INFO    MainThread:13227 [wandb_setup.py:_flush():80] Loading settings from /home/clouduser/Code/Github/unified_world_model/wandb/settings
-2026-01-02 21:43:04,973 INFO    MainThread:13227 [wandb_setup.py:_flush():80] Loading settings from environment variables
-2026-01-02 21:43:04,973 INFO    MainThread:13227 [wandb_init.py:setup_run_log_directory():714] Logging user logs to /dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_mse_only/wandb/offline-run-20260102_214304-vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0/logs/debug.log
-2026-01-02 21:43:04,973 INFO    MainThread:13227 [wandb_init.py:setup_run_log_directory():715] Logging internal logs to /dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_mse_only/wandb/offline-run-20260102_214304-vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0/logs/debug-internal.log
-2026-01-02 21:43:04,973 INFO    MainThread:13227 [wandb_init.py:init():841] calling init triggers
-2026-01-02 21:43:04,973 INFO    MainThread:13227 [wandb_init.py:init():846] wandb.init called with sweep_config: {}
+2026-01-04 09:32:00,479 INFO    MainThread:13221 [wandb_setup.py:_flush():80] Current SDK version is 0.23.1
+2026-01-04 09:32:00,479 INFO    MainThread:13221 [wandb_setup.py:_flush():80] Configure stats pid to 13221
+2026-01-04 09:32:00,479 INFO    MainThread:13221 [wandb_setup.py:_flush():80] Loading settings from /home/clouduser/.config/wandb/settings
+2026-01-04 09:32:00,479 INFO    MainThread:13221 [wandb_setup.py:_flush():80] Loading settings from /home/clouduser/Code/Github/unified_world_model/wandb/settings
+2026-01-04 09:32:00,479 INFO    MainThread:13221 [wandb_setup.py:_flush():80] Loading settings from environment variables
+2026-01-04 09:32:00,479 INFO    MainThread:13221 [wandb_init.py:setup_run_log_directory():714] Logging user logs to /dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_mse_only/wandb/offline-run-20260104_093200-vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0/logs/debug.log
+2026-01-04 09:32:00,479 INFO    MainThread:13221 [wandb_init.py:setup_run_log_directory():715] Logging internal logs to /dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_mse_only/wandb/offline-run-20260104_093200-vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0/logs/debug-internal.log
+2026-01-04 09:32:00,479 INFO    MainThread:13221 [wandb_init.py:init():841] calling init triggers
+2026-01-04 09:32:00,480 INFO    MainThread:13221 [wandb_init.py:init():846] wandb.init called with sweep_config: {}
 config: {'_wandb': {}}
-2026-01-02 21:43:04,973 INFO    MainThread:13227 [wandb_init.py:init():889] starting backend
-2026-01-02 21:43:05,456 INFO    MainThread:13227 [wandb_init.py:init():892] sending inform_init request
-2026-01-02 21:43:05,464 INFO    MainThread:13227 [wandb_init.py:init():900] backend started and connected
-2026-01-02 21:43:05,466 INFO    MainThread:13227 [wandb_init.py:init():970] updated telemetry
-2026-01-02 21:43:05,474 INFO    MainThread:13227 [wandb_init.py:init():994] communicating run to backend with 90.0 second timeout
-2026-01-02 21:43:06,064 INFO    MainThread:13227 [wandb_init.py:init():1041] starting run threads in backend
-2026-01-02 21:43:06,498 INFO    MainThread:13227 [wandb_run.py:_console_start():2521] atexit reg
-2026-01-02 21:43:06,498 INFO    MainThread:13227 [wandb_run.py:_redirect():2369] redirect: wrap_raw
-2026-01-02 21:43:06,498 INFO    MainThread:13227 [wandb_run.py:_redirect():2438] Wrapping output streams.
-2026-01-02 21:43:06,498 INFO    MainThread:13227 [wandb_run.py:_redirect():2461] Redirects installed.
-2026-01-02 21:43:06,502 INFO    MainThread:13227 [wandb_init.py:init():1081] run started, returning control to user process
-2026-01-02 21:43:06,504 INFO    MainThread:13227 [wandb_run.py:_config_callback():1396] config_cb None None {'visual_gen': True, 'visual_und': True, 'results_dir': 'results', 'checkpoint_dir': '/dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_mse_only', 'wandb_project': 'bagel', 'wandb_name': 'vlm_gym_jigsaw_one_img_lr2e_5_mse_only', 'wandb_runid': '0', 'wandb_resume': 'allow', 'wandb_offline': True, 'wandb_dir': '/dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_mse_only', 'global_seed': 4396, 'auto_resume': False, 'resume_from': '/home/clouduser/Code/Models/BAGEL-7B-MoT', 'resume_model_only': True, 'finetune_from_ema': True, 'finetune_from_hf': True, 'log_every': 1, 'save_every': 2500, 'total_steps': 5000, 'warmup_steps': 300, 'lr_scheduler': 'cosine', 'lr': 2e-05, 'min_lr': 1e-07, 'beta1': 0.9, 'beta2': 0.95, 'eps': 1e-15, 'ema': 0.993, 'max_grad_norm': 1.0, 'timestep_shift': 1.0, 'mse_weight': 1.0, 'ce_weight': 1.0, 'ce_loss_reweighting': False, 'expected_num_tokens': 20000, 'num_replicate': 1, 'num_shard': 8, 'sharding_strategy': 'HYBRID_SHARD', 'backward_prefetch': 'BACKWARD_PRE', 'cpu_offload': False, 'freeze_llm': False, 'freeze_vit': False, 'freeze_vae': True, 'freeze_und': False, 'copy_init_moe': True, 'use_flex': False, 'eval_every': 500, 'num_eval_batches': 20, 'use_ema_for_eval': True, 'viz_every': 500, 'viz_n': 8, 'viz_outdir': 'results/viz', 'eval_dataset_config_file': './data/configs/vlm_gym_jigsaw_train_mseloss_only.yaml', 'viz_dataset_config_file': './data/configs/vlm_gym_jigsaw_train_mseloss_only.yaml', 'save_ema_only': True, 'save_optimizer': False}
-2026-01-02 21:43:06,505 INFO    MainThread:13227 [wandb_run.py:_config_callback():1396] config_cb None None {'model_path': '/home/clouduser/Code/Models/BAGEL-7B-MoT', 'llm_path': 'hf/Qwen2.5-0.5B-Instruct/', 'llm_qk_norm': True, 'tie_word_embeddings': False, 'layer_module': 'Qwen2MoTDecoderLayer', 'vae_path': 'flux/vae/ae.safetensors', 'vit_path': 'hf/siglip-so400m-14-980-flash-attn2-navit/', 'max_latent_size': 64, 'latent_patch_size': 2, 'vit_patch_size': 14, 'vit_max_num_patch_per_side': 70, 'connector_act': 'gelu_pytorch_tanh', 'interpolate_pos': False, 'vit_select_layer': -2, 'vit_rope': False, 'text_cond_dropout_prob': 0.0, 'vae_cond_dropout_prob': 0.0, 'vit_cond_dropout_prob': 0.0}
-2026-01-02 21:43:06,505 INFO    MainThread:13227 [wandb_run.py:_config_callback():1396] config_cb None None {'dataset_config_file': './data/configs/vlm_gym_jigsaw_train_mseloss_only.yaml', 'train_data_dir': '/home/clouduser/Code/data/gym/jigsaw-swap_v5/train/', 'train_jsonl_path': '/home/clouduser/Code/data/gym/jigsaw-swap_v5/train/', 'eval_data_dir': '/home/clouduser/Code/data/gym/jigsaw-swap_v5/val/', 'eval_jsonl_path': '/home/clouduser/Code/data/gym/jigsaw-swap_v5/val/', 'inference_hash_file': '/home/clouduser/Code/Github/launch_new/hashes_test_set_v10.json', 'prefetch_factor': 2, 'num_workers': 1, 'max_num_tokens_per_sample': 20000, 'max_num_tokens': 20000, 'prefer_buffer_before': 16384, 'max_buffer_size': 50, 'data_seed': 42}
-2026-01-03 17:12:43,959 INFO    MainThread:13227 [wandb_run.py:_finish():2287] finishing run bagel/vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0
-2026-01-03 17:12:43,960 INFO    MainThread:13227 [wandb_run.py:_atexit_cleanup():2486] got exitcode: 0
-2026-01-03 17:12:43,960 INFO    MainThread:13227 [wandb_run.py:_restore():2468] restore
-2026-01-03 17:12:43,960 INFO    MainThread:13227 [wandb_run.py:_restore():2474] restore done
+2026-01-04 09:32:00,480 INFO    MainThread:13221 [wandb_init.py:init():889] starting backend
+2026-01-04 09:32:00,939 INFO    MainThread:13221 [wandb_init.py:init():892] sending inform_init request
+2026-01-04 09:32:00,947 INFO    MainThread:13221 [wandb_init.py:init():900] backend started and connected
+2026-01-04 09:32:00,949 INFO    MainThread:13221 [wandb_init.py:init():970] updated telemetry
+2026-01-04 09:32:00,956 INFO    MainThread:13221 [wandb_init.py:init():994] communicating run to backend with 90.0 second timeout
+2026-01-04 09:32:01,446 INFO    MainThread:13221 [wandb_init.py:init():1041] starting run threads in backend
+2026-01-04 09:32:01,964 INFO    MainThread:13221 [wandb_run.py:_console_start():2521] atexit reg
+2026-01-04 09:32:01,964 INFO    MainThread:13221 [wandb_run.py:_redirect():2369] redirect: wrap_raw
+2026-01-04 09:32:01,964 INFO    MainThread:13221 [wandb_run.py:_redirect():2438] Wrapping output streams.
+2026-01-04 09:32:01,964 INFO    MainThread:13221 [wandb_run.py:_redirect():2461] Redirects installed.
+2026-01-04 09:32:01,968 INFO    MainThread:13221 [wandb_init.py:init():1081] run started, returning control to user process
+2026-01-04 09:32:01,969 INFO    MainThread:13221 [wandb_run.py:_config_callback():1396] config_cb None None {'visual_gen': True, 'visual_und': True, 'results_dir': 'results', 'checkpoint_dir': '/dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_mse_only', 'wandb_project': 'bagel', 'wandb_name': 'vlm_gym_jigsaw_one_img_lr2e_5_mse_only', 'wandb_runid': '0', 'wandb_resume': 'allow', 'wandb_offline': True, 'wandb_dir': '/dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_mse_only', 'global_seed': 4396, 'auto_resume': False, 'resume_from': '/home/clouduser/Code/Models/BAGEL-7B-MoT', 'resume_model_only': True, 'finetune_from_ema': True, 'finetune_from_hf': True, 'log_every': 1, 'save_every': 2500, 'total_steps': 5000, 'warmup_steps': 300, 'lr_scheduler': 'cosine', 'lr': 2e-05, 'min_lr': 1e-07, 'beta1': 0.9, 'beta2': 0.95, 'eps': 1e-15, 'ema': 0.993, 'max_grad_norm': 1.0, 'timestep_shift': 1.0, 'mse_weight': 1.0, 'ce_weight': 1.0, 'ce_loss_reweighting': False, 'expected_num_tokens': 20000, 'num_replicate': 1, 'num_shard': 8, 'sharding_strategy': 'HYBRID_SHARD', 'backward_prefetch': 'BACKWARD_PRE', 'cpu_offload': False, 'freeze_llm': False, 'freeze_vit': False, 'freeze_vae': True, 'freeze_und': False, 'copy_init_moe': True, 'use_flex': False, 'eval_every': 500, 'num_eval_batches': 20, 'use_ema_for_eval': True, 'eval_log_dir': None, 'eval_run_tag': '', 'viz_every': 500, 'viz_n': 8, 'viz_outdir': 'results/viz', 'eval_dataset_config_file': './data/configs/vlm_gym_jigsaw_train_mseloss_only.yaml', 'viz_dataset_config_file': './data/configs/vlm_gym_jigsaw_train_mseloss_only.yaml', 'eval_print_n': 3, 'save_ema_only': True, 'save_optimizer': False}
+2026-01-04 09:32:01,970 INFO    MainThread:13221 [wandb_run.py:_config_callback():1396] config_cb None None {'model_path': '/home/clouduser/Code/Models/BAGEL-7B-MoT', 'llm_path': 'hf/Qwen2.5-0.5B-Instruct/', 'llm_qk_norm': True, 'tie_word_embeddings': False, 'layer_module': 'Qwen2MoTDecoderLayer', 'vae_path': 'flux/vae/ae.safetensors', 'vit_path': 'hf/siglip-so400m-14-980-flash-attn2-navit/', 'max_latent_size': 64, 'latent_patch_size': 2, 'vit_patch_size': 14, 'vit_max_num_patch_per_side': 70, 'connector_act': 'gelu_pytorch_tanh', 'interpolate_pos': False, 'vit_select_layer': -2, 'vit_rope': False, 'text_cond_dropout_prob': 0.0, 'vae_cond_dropout_prob': 0.0, 'vit_cond_dropout_prob': 0.0}
+2026-01-04 09:32:01,971 INFO    MainThread:13221 [wandb_run.py:_config_callback():1396] config_cb None None {'dataset_config_file': './data/configs/vlm_gym_jigsaw_train_mseloss_only.yaml', 'train_data_dir': '/home/clouduser/Code/data/gym/jigsaw-swap_v5/train/', 'train_jsonl_path': '/home/clouduser/Code/data/gym/jigsaw-swap_v5/train/', 'eval_data_dir': '/home/clouduser/Code/data/gym/jigsaw-swap_v5/val/', 'eval_jsonl_path': '/home/clouduser/Code/data/gym/jigsaw-swap_v5/val/', 'inference_hash_file': '/home/clouduser/Code/Github/launch_new/hashes_test_set_v10.json', 'prefetch_factor': 2, 'num_workers': 1, 'max_num_tokens_per_sample': 20000, 'max_num_tokens': 20000, 'prefer_buffer_before': 16384, 'max_buffer_size': 50, 'data_seed': 42}

checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_mse_only/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_mse_only/wandb/offline-run-20251230_194213-vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0/files/output.log

@@ -1,176 +1,3 @@
-[2025-12-30 23:03:51] (step=0000844) Train Loss mse: 0.0270, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-FullyShardedDataParallel(
-  (_fsdp_wrapped_module): Bagel(
-    (language_model): Qwen2ForCausalLM(
-      (model): Qwen2Model(
-        (embed_tokens): Embedding(152064, 3584)
-        (layers): ModuleList(
-          (0-27): 28 x FullyShardedDataParallel(
-            (_fsdp_wrapped_module): CheckpointWrapper(
-              (_checkpoint_wrapped_module): Qwen2MoTDecoderLayer(
-                (self_attn): PackedAttentionMoT(
-                  (q_proj): Linear(in_features=3584, out_features=3584, bias=True)
-                  (k_proj): Linear(in_features=3584, out_features=512, bias=True)
-                  (v_proj): Linear(in_features=3584, out_features=512, bias=True)
-                  (o_proj): Linear(in_features=3584, out_features=3584, bias=False)
-                  (q_norm): Qwen2RMSNorm((128,), eps=1e-06)
-                  (k_norm): Qwen2RMSNorm((128,), eps=1e-06)
-                  (q_norm_moe_gen): Qwen2RMSNorm((128,), eps=1e-06)
-                  (k_norm_moe_gen): Qwen2RMSNorm((128,), eps=1e-06)
-                  (q_proj_moe_gen): Linear(in_features=3584, out_features=3584, bias=True)
-                  (k_proj_moe_gen): Linear(in_features=3584, out_features=512, bias=True)
-                  (v_proj_moe_gen): Linear(in_features=3584, out_features=512, bias=True)
-                  (o_proj_moe_gen): Linear(in_features=3584, out_features=3584, bias=False)
-                )
-                (mlp): Qwen2MLP(
-                  (gate_proj): Linear(in_features=3584, out_features=18944, bias=False)
-                  (up_proj): Linear(in_features=3584, out_features=18944, bias=False)
-                  (down_proj): Linear(in_features=18944, out_features=3584, bias=False)
-                  (act_fn): SiLU()
-                )
-                (mlp_moe_gen): Qwen2MLP(
-                  (gate_proj): Linear(in_features=3584, out_features=18944, bias=False)
-                  (up_proj): Linear(in_features=3584, out_features=18944, bias=False)
-                  (down_proj): Linear(in_features=18944, out_features=3584, bias=False)
-                  (act_fn): SiLU()
-                )
-                (input_layernorm): Qwen2RMSNorm((3584,), eps=1e-06)
-                (input_layernorm_moe_gen): Qwen2RMSNorm((3584,), eps=1e-06)
-                (post_attention_layernorm): Qwen2RMSNorm((3584,), eps=1e-06)
-                (post_attention_layernorm_moe_gen): Qwen2RMSNorm((3584,), eps=1e-06)
-              )
-            )
-          )
-        )
-        (norm): Qwen2RMSNorm((3584,), eps=1e-06)
-        (norm_moe_gen): Qwen2RMSNorm((3584,), eps=1e-06)
-        (rotary_emb): Qwen2RotaryEmbedding()
-      )
-      (lm_head): Linear(in_features=3584, out_features=152064, bias=False)
-    )
-    (time_embedder): FullyShardedDataParallel(
-      (_fsdp_wrapped_module): TimestepEmbedder(
-        (mlp): Sequential(
-          (0): Linear(in_features=256, out_features=3584, bias=True)
-          (1): SiLU()
-          (2): Linear(in_features=3584, out_features=3584, bias=True)
-        )
-      )
-    )
-    (vae2llm): Linear(in_features=64, out_features=3584, bias=True)
-    (llm2vae): Linear(in_features=3584, out_features=64, bias=True)
-    (latent_pos_embed): FullyShardedDataParallel(
-      (_fsdp_wrapped_module): PositionEmbedding()
-    )
-    (vit_model): SiglipVisionModel(
-      (vision_model): FullyShardedDataParallel(
-        (_fsdp_wrapped_module): SiglipVisionTransformer(
-          (embeddings): SiglipVisionEmbeddings(
-            (position_embedding): Embedding(4900, 1152)
-            (patch_embedding): Linear(in_features=588, out_features=1152, bias=True)
-          )
-          (encoder): SiglipEncoder(
-            (layers): ModuleList(
-              (0-25): 26 x FullyShardedDataParallel(
-                (_fsdp_wrapped_module): CheckpointWrapper(
-                  (_checkpoint_wrapped_module): SiglipEncoderLayer(
-                    (self_attn): SiglipFlashAttention2(
-                      (k_proj): Linear(in_features=1152, out_features=1152, bias=True)
-                      (v_proj): Linear(in_features=1152, out_features=1152, bias=True)
-                      (q_proj): Linear(in_features=1152, out_features=1152, bias=True)
-                      (out_proj): Linear(in_features=1152, out_features=1152, bias=True)
-                    )
-                    (layer_norm1): LayerNorm((1152,), eps=1e-06, elementwise_affine=True)
-                    (mlp): SiglipMLP(
-                      (activation_fn): PytorchGELUTanh()
-                      (fc1): Linear(in_features=1152, out_features=4304, bias=True)
-                      (fc2): Linear(in_features=4304, out_features=1152, bias=True)
-                    )
-                    (layer_norm2): LayerNorm((1152,), eps=1e-06, elementwise_affine=True)
-                  )
-                )
-              )
-            )
-          )
-          (post_layernorm): LayerNorm((1152,), eps=1e-06, elementwise_affine=True)
-        )
-      )
-    )
-    (connector): FullyShardedDataParallel(
-      (_fsdp_wrapped_module): CheckpointWrapper(
-        (_checkpoint_wrapped_module): MLPconnector(
-          (activation_fn): PytorchGELUTanh()
-          (fc1): Linear(in_features=1152, out_features=3584, bias=True)
-          (fc2): Linear(in_features=3584, out_features=3584, bias=True)
-        )
-      )
-    )
-    (vit_pos_embed): FullyShardedDataParallel(
-      (_fsdp_wrapped_module): PositionEmbedding()
-    )
-  )
-)
-_flat_param True
-language_model.model.layers.0._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.1._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.2._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.3._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.4._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.5._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.6._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.7._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.8._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.9._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.10._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.11._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.12._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.13._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.14._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.15._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.16._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.17._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.18._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.19._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.20._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.21._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.22._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.23._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.24._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.25._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.26._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.27._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-time_embedder._fsdp_wrapped_module._flat_param True
-latent_pos_embed._fsdp_wrapped_module._flat_param False
-vit_model.vision_model._fsdp_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.0._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.1._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.2._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.3._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.4._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.5._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.6._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.7._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.8._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.9._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.10._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.11._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.12._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.13._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.14._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.15._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.16._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.17._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.18._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.19._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.20._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.21._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.22._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.23._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.24._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.25._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-connector._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_pos_embed._fsdp_wrapped_module._flat_param False
-Preparing Dataset vlm_gym_jigsaw_mse_loss_only/vlm_gym_jigsaw_train
 wandb: Detected [huggingface_hub.inference] in use.
 wandb: Use W&B Weave for improved LLM call tracing. Install Weave with `pip install weave` then add `import weave` to the top of your script.
 wandb: For more information, check out the docs at: https://weave-docs.wandb.ai/
@@ -1007,6 +834,179 @@ wandb: For more information, check out the docs at: https://weave-docs.wandb.ai/
 [2025-12-30 22:59:02] (step=0000823) Train Loss mse: 0.0429, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
 [2025-12-30 22:59:18] (step=0000824) Train Loss mse: 0.0383, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
 [2025-12-30 22:59:30] (step=0000825) Train Loss mse: 0.0483, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
+FullyShardedDataParallel(
+  (_fsdp_wrapped_module): Bagel(
+    (language_model): Qwen2ForCausalLM(
+      (model): Qwen2Model(
+        (embed_tokens): Embedding(152064, 3584)
+        (layers): ModuleList(
+          (0-27): 28 x FullyShardedDataParallel(
+            (_fsdp_wrapped_module): CheckpointWrapper(
+              (_checkpoint_wrapped_module): Qwen2MoTDecoderLayer(
+                (self_attn): PackedAttentionMoT(
+                  (q_proj): Linear(in_features=3584, out_features=3584, bias=True)
+                  (k_proj): Linear(in_features=3584, out_features=512, bias=True)
+                  (v_proj): Linear(in_features=3584, out_features=512, bias=True)
+                  (o_proj): Linear(in_features=3584, out_features=3584, bias=False)
+                  (q_norm): Qwen2RMSNorm((128,), eps=1e-06)
+                  (k_norm): Qwen2RMSNorm((128,), eps=1e-06)
+                  (q_norm_moe_gen): Qwen2RMSNorm((128,), eps=1e-06)
+                  (k_norm_moe_gen): Qwen2RMSNorm((128,), eps=1e-06)
+                  (q_proj_moe_gen): Linear(in_features=3584, out_features=3584, bias=True)
+                  (k_proj_moe_gen): Linear(in_features=3584, out_features=512, bias=True)
+                  (v_proj_moe_gen): Linear(in_features=3584, out_features=512, bias=True)
+                  (o_proj_moe_gen): Linear(in_features=3584, out_features=3584, bias=False)
+                )
+                (mlp): Qwen2MLP(
+                  (gate_proj): Linear(in_features=3584, out_features=18944, bias=False)
+                  (up_proj): Linear(in_features=3584, out_features=18944, bias=False)
+                  (down_proj): Linear(in_features=18944, out_features=3584, bias=False)
+                  (act_fn): SiLU()
+                )
+                (mlp_moe_gen): Qwen2MLP(
+                  (gate_proj): Linear(in_features=3584, out_features=18944, bias=False)
+                  (up_proj): Linear(in_features=3584, out_features=18944, bias=False)
+                  (down_proj): Linear(in_features=18944, out_features=3584, bias=False)
+                  (act_fn): SiLU()
+                )
+                (input_layernorm): Qwen2RMSNorm((3584,), eps=1e-06)
+                (input_layernorm_moe_gen): Qwen2RMSNorm((3584,), eps=1e-06)
+                (post_attention_layernorm): Qwen2RMSNorm((3584,), eps=1e-06)
+                (post_attention_layernorm_moe_gen): Qwen2RMSNorm((3584,), eps=1e-06)
+              )
+            )
+          )
+        )
+        (norm): Qwen2RMSNorm((3584,), eps=1e-06)
+        (norm_moe_gen): Qwen2RMSNorm((3584,), eps=1e-06)
+        (rotary_emb): Qwen2RotaryEmbedding()
+      )
+      (lm_head): Linear(in_features=3584, out_features=152064, bias=False)
+    )
+    (time_embedder): FullyShardedDataParallel(
+      (_fsdp_wrapped_module): TimestepEmbedder(
+        (mlp): Sequential(
+          (0): Linear(in_features=256, out_features=3584, bias=True)
+          (1): SiLU()
+          (2): Linear(in_features=3584, out_features=3584, bias=True)
+        )
+      )
+    )
+    (vae2llm): Linear(in_features=64, out_features=3584, bias=True)
+    (llm2vae): Linear(in_features=3584, out_features=64, bias=True)
+    (latent_pos_embed): FullyShardedDataParallel(
+      (_fsdp_wrapped_module): PositionEmbedding()
+    )
+    (vit_model): SiglipVisionModel(
+      (vision_model): FullyShardedDataParallel(
+        (_fsdp_wrapped_module): SiglipVisionTransformer(
+          (embeddings): SiglipVisionEmbeddings(
+            (position_embedding): Embedding(4900, 1152)
+            (patch_embedding): Linear(in_features=588, out_features=1152, bias=True)
+          )
+          (encoder): SiglipEncoder(
+            (layers): ModuleList(
+              (0-25): 26 x FullyShardedDataParallel(
+                (_fsdp_wrapped_module): CheckpointWrapper(
+                  (_checkpoint_wrapped_module): SiglipEncoderLayer(
+                    (self_attn): SiglipFlashAttention2(
+                      (k_proj): Linear(in_features=1152, out_features=1152, bias=True)
+                      (v_proj): Linear(in_features=1152, out_features=1152, bias=True)
+                      (q_proj): Linear(in_features=1152, out_features=1152, bias=True)
+                      (out_proj): Linear(in_features=1152, out_features=1152, bias=True)
+                    )
+                    (layer_norm1): LayerNorm((1152,), eps=1e-06, elementwise_affine=True)
+                    (mlp): SiglipMLP(
+                      (activation_fn): PytorchGELUTanh()
+                      (fc1): Linear(in_features=1152, out_features=4304, bias=True)
+                      (fc2): Linear(in_features=4304, out_features=1152, bias=True)
+                    )
+                    (layer_norm2): LayerNorm((1152,), eps=1e-06, elementwise_affine=True)
+                  )
+                )
+              )
+            )
+          )
+          (post_layernorm): LayerNorm((1152,), eps=1e-06, elementwise_affine=True)
+        )
+      )
+    )
+    (connector): FullyShardedDataParallel(
+      (_fsdp_wrapped_module): CheckpointWrapper(
+        (_checkpoint_wrapped_module): MLPconnector(
+          (activation_fn): PytorchGELUTanh()
+          (fc1): Linear(in_features=1152, out_features=3584, bias=True)
+          (fc2): Linear(in_features=3584, out_features=3584, bias=True)
+        )
+      )
+    )
+    (vit_pos_embed): FullyShardedDataParallel(
+      (_fsdp_wrapped_module): PositionEmbedding()
+    )
+  )
+)
+_flat_param True
+language_model.model.layers.0._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.1._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.2._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.3._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.4._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.5._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.6._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.7._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.8._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.9._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.10._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.11._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.12._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.13._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.14._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.15._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.16._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.17._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.18._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.19._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.20._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.21._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.22._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.23._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.24._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.25._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.26._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.27._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+time_embedder._fsdp_wrapped_module._flat_param True
+latent_pos_embed._fsdp_wrapped_module._flat_param False
+vit_model.vision_model._fsdp_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.0._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.1._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.2._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.3._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.4._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.5._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.6._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.7._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.8._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.9._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.10._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.11._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.12._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.13._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.14._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.15._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.16._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.17._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.18._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.19._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.20._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.21._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.22._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.23._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.24._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.25._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+connector._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_pos_embed._fsdp_wrapped_module._flat_param False
+Preparing Dataset vlm_gym_jigsaw_mse_loss_only/vlm_gym_jigsaw_train
+Preparing Dataset vlm_gym_jigsaw_mse_loss_only/vlm_gym_jigsaw_train
 [2025-12-30 22:59:42] (step=0000826) Train Loss mse: 0.0364, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
 [2025-12-30 22:59:55] (step=0000827) Train Loss mse: 0.0403, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
 [2025-12-30 23:00:08] (step=0000828) Train Loss mse: 0.0478, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
@@ -1025,7 +1025,7 @@ wandb: For more information, check out the docs at: https://weave-docs.wandb.ai/
 [2025-12-30 23:03:06] (step=0000841) Train Loss mse: 0.0491, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
 [2025-12-30 23:03:19] (step=0000842) Train Loss mse: 0.0546, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
 [2025-12-30 23:03:35] (step=0000843) Train Loss mse: 0.0530, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-Preparing Dataset vlm_gym_jigsaw_mse_loss_only/vlm_gym_jigsaw_train
+[2025-12-30 23:03:51] (step=0000844) Train Loss mse: 0.0270, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
 [2025-12-30 23:04:04] (step=0000845) Train Loss mse: 0.0425, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
 [2025-12-30 23:04:17] (step=0000846) Train Loss mse: 0.0424, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
 [2025-12-30 23:04:33] (step=0000847) Train Loss mse: 0.0337, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
@@ -3132,29 +3132,7 @@ Preparing Dataset vlm_gym_jigsaw_mse_loss_only/vlm_gym_jigsaw_train
 [2025-12-31 07:13:08] (step=0002945) Train Loss mse: 0.0254, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
 [2025-12-31 07:13:21] (step=0002946) Train Loss mse: 0.0386, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
 [2025-12-31 07:13:34] (step=0002947) Train Loss mse: 0.0290, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2025-12-31 07:13:48] (step=0002948) Train Loss mse: 0.0196, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2025-12-31 07:14:03] (step=0002949) Train Loss mse: 0.0272, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2025-12-31 07:14:19] (step=0002950) Train Loss mse: 0.0222, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2025-12-31 07:14:29] (step=0002951) Train Loss mse: 0.0206, Train Loss ce: 0.0000, Train Steps/Sec: 0.11,
-[2025-12-31 07:14:41] (step=0002952) Train Loss mse: 0.0373, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2025-12-31 07:14:54] (step=0002953) Train Loss mse: 0.0368, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2025-12-31 07:15:08] (step=0002954) Train Loss mse: 0.0234, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2025-12-31 07:15:24] (step=0002955) Train Loss mse: 0.0337, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2025-12-31 07:15:37] (step=0002956) Train Loss mse: 0.0295, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2025-12-31 07:15:53] (step=0002957) Train Loss mse: 0.0307, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2025-12-31 07:16:06] (step=0002958) Train Loss mse: 0.0301, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2025-12-31 07:16:17] (step=0002959) Train Loss mse: 0.0331, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
-[2025-12-31 07:16:33] (step=0002960) Train Loss mse: 0.0245, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2025-12-31 07:16:47] (step=0002961) Train Loss mse: 0.0196, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2025-12-31 07:17:00] (step=0002962) Train Loss mse: 0.0269, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2025-12-31 07:17:14] (step=0002963) Train Loss mse: 0.0276, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2025-12-31 07:17:30] (step=0002964) Train Loss mse: 0.0233, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2025-12-31 07:17:46] (step=0002965) Train Loss mse: 0.0461, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2025-12-31 07:17:57] (step=0002966) Train Loss mse: 0.0241, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
-[2025-12-31 07:18:10] (step=0002967) Train Loss mse: 0.0296, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2025-12-31 07:18:24] (step=0002968) Train Loss mse: 0.0290, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2025-12-31 07:18:40] (step=0002969) Train Loss mse: 0.0319, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2025-12-31 07:18:56] (step=0002970) Train Loss mse: 0.0254, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2025-12-31 07:13:48
 [2025-12-31 07:19:10] (step=0002971) Train Loss mse: 0.0283, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
 [2025-12-31 07:19:22] (step=0002972) Train Loss mse: 0.0399, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
 [2025-12-31 07:19:38] (step=0002973) Train Loss mse: 0.0229, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,

checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_mse_only/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_mse_only/wandb/offline-run-20260102_214304-vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0/files/output.log

@@ -747,6 +747,121 @@ wandb: For more information, check out the docs at: https://weave-docs.wandb.ai/
 [2026-01-03 00:39:54] (step=0000736) Train Loss mse: 0.0265, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
 [2026-01-03 00:40:11] (step=0000737) Train Loss mse: 0.0442, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
 [2026-01-03 00:40:24] (step=0000738) Train Loss mse: 0.0433, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 00:40:37] (step=0000739) Train Loss mse: 0.0488, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 00:40:50] (step=0000740) Train Loss mse: 0.0456, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 00:41:04] (step=0000741) Train Loss mse: 0.0518, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 00:41:16] (step=0000742) Train Loss mse: 0.0441, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 00:41:27] (step=0000743) Train Loss mse: 0.0542, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
+[2026-01-03 00:41:41] (step=0000744) Train Loss mse: 0.0310, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 00:41:53] (step=0000745) Train Loss mse: 0.0629, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 00:42:09] (step=0000746) Train Loss mse: 0.0451, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 00:42:23] (step=0000747) Train Loss mse: 0.0331, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 00:42:36] (step=0000748) Train Loss mse: 0.0387, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 00:42:47] (step=0000749) Train Loss mse: 0.0353, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
+[2026-01-03 00:42:59] (step=0000750) Train Loss mse: 0.0455, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
+[2026-01-03 00:43:12] (step=0000751) Train Loss mse: 0.0436, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 00:43:25] (step=0000752) Train Loss mse: 0.0441, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 00:43:41] (step=0000753) Train Loss mse: 0.0404, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 00:43:51] (step=0000754) Train Loss mse: 0.0379, Train Loss ce: 0.0000, Train Steps/Sec: 0.10,
+[2026-01-03 00:44:04] (step=0000755) Train Loss mse: 0.0340, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 00:44:20] (step=0000756) Train Loss mse: 0.0289, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 00:44:36] (step=0000757) Train Loss mse: 0.0436, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 00:44:52] (step=0000758) Train Loss mse: 0.0460, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 00:45:05] (step=0000759) Train Loss mse: 0.0317, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 00:45:22] (step=0000760) Train Loss mse: 0.0320, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 00:45:35] (step=0000761) Train Loss mse: 0.0383, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 00:45:46] (step=0000762) Train Loss mse: 0.0434, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
+[2026-01-03 00:46:03] (step=0000763) Train Loss mse: 0.0465, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 00:46:14] (step=0000764) Train Loss mse: 0.0320, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
+[2026-01-03 00:46:27] (step=0000765) Train Loss mse: 0.0423, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 00:46:40] (step=0000766) Train Loss mse: 0.0428, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 00:46:52] (step=0000767) Train Loss mse: 0.0392, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 00:47:06] (step=0000768) Train Loss mse: 0.0506, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 00:47:22] (step=0000769) Train Loss mse: 0.0424, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 00:47:34] (step=0000770) Train Loss mse: 0.0405, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 00:47:49] (step=0000771) Train Loss mse: 0.0442, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 00:48:02] (step=0000772) Train Loss mse: 0.0469, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 00:48:15] (step=0000773) Train Loss mse: 0.0456, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 00:48:28] (step=0000774) Train Loss mse: 0.0431, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 00:48:44] (step=0000775) Train Loss mse: 0.0369, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 00:49:00] (step=0000776) Train Loss mse: 0.0357, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 00:49:14] (step=0000777) Train Loss mse: 0.0466, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 00:49:30] (step=0000778) Train Loss mse: 0.0367, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 00:49:46] (step=0000779) Train Loss mse: 0.0452, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 00:49:56] (step=0000780) Train Loss mse: 0.0635, Train Loss ce: 0.0000, Train Steps/Sec: 0.10,
+[2026-01-03 00:50:08] (step=0000781) Train Loss mse: 0.0540, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 00:50:24] (step=0000782) Train Loss mse: 0.0468, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 00:50:36] (step=0000783) Train Loss mse: 0.0439, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 00:50:48] (step=0000784) Train Loss mse: 0.0506, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 00:51:00] (step=0000785) Train Loss mse: 0.0376, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 00:51:16] (step=0000786) Train Loss mse: 0.0326, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 00:51:29] (step=0000787) Train Loss mse: 0.0364, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 00:51:41] (step=0000788) Train Loss mse: 0.0292, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 00:51:57] (step=0000789) Train Loss mse: 0.0470, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 00:52:13] (step=0000790) Train Loss mse: 0.0311, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 00:52:27] (step=0000791) Train Loss mse: 0.0399, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 00:52:43] (step=0000792) Train Loss mse: 0.0250, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 00:52:55] (step=0000793) Train Loss mse: 0.0513, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 00:53:06] (step=0000794) Train Loss mse: 0.0426, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
+[2026-01-03 00:53:20] (step=0000795) Train Loss mse: 0.0637, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 00:53:33] (step=0000796) Train Loss mse: 0.0504, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 00:53:48] (step=0000797) Train Loss mse: 0.0467, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 00:54:00] (step=0000798) Train Loss mse: 0.0493, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
+[2026-01-03 00:54:16] (step=0000799) Train Loss mse: 0.0400, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 00:54:32] (step=0000800) Train Loss mse: 0.0519, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 00:54:44] (step=0000801) Train Loss mse: 0.0510, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 00:54:58] (step=0000802) Train Loss mse: 0.0425, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 00:55:09] (step=0000803) Train Loss mse: 0.0396, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
+[2026-01-03 00:55:19] (step=0000804) Train Loss mse: 0.0541, Train Loss ce: 0.0000, Train Steps/Sec: 0.11,
+[2026-01-03 00:55:32] (step=0000805) Train Loss mse: 0.0452, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 00:55:43] (step=0000806) Train Loss mse: 0.0527, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
+[2026-01-03 00:55:56] (step=0000807) Train Loss mse: 0.0441, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 00:56:13] (step=0000808) Train Loss mse: 0.0269, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 00:56:26] (step=0000809) Train Loss mse: 0.0419, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 00:56:39] (step=0000810) Train Loss mse: 0.0485, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 00:56:53] (step=0000811) Train Loss mse: 0.0351, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 00:57:09] (step=0000812) Train Loss mse: 0.0273, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 00:57:25] (step=0000813) Train Loss mse: 0.0412, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 00:57:37] (step=0000814) Train Loss mse: 0.0527, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 00:57:53] (step=0000815) Train Loss mse: 0.0316, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 00:58:06] (step=0000816) Train Loss mse: 0.0294, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 00:58:23] (step=0000817) Train Loss mse: 0.0337, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 00:58:39] (step=0000818) Train Loss mse: 0.0415, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 00:58:52] (step=0000819) Train Loss mse: 0.0356, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 00:59:05] (step=0000820) Train Loss mse: 0.0428, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 00:59:19] (step=0000821) Train Loss mse: 0.0456, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 00:59:33] (step=0000822) Train Loss mse: 0.0497, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 00:59:46] (step=0000823) Train Loss mse: 0.0428, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 01:00:02] (step=0000824) Train Loss mse: 0.0385, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 01:00:13] (step=0000825) Train Loss mse: 0.0499, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
+[2026-01-03 01:00:25] (step=0000826) Train Loss mse: 0.0360, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 01:00:39] (step=0000827) Train Loss mse: 0.0415, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 01:00:52] (step=0000828) Train Loss mse: 0.0479, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 01:01:05] (step=0000829) Train Loss mse: 0.0361, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 01:01:17] (step=0000830) Train Loss mse: 0.0509, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
+[2026-01-03 01:01:30] (step=0000831) Train Loss mse: 0.0425, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 01:01:43] (step=0000832) Train Loss mse: 0.0365, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 01:01:59] (step=0000833) Train Loss mse: 0.0427, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 01:02:11] (step=0000834) Train Loss mse: 0.0505, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 01:02:24] (step=0000835) Train Loss mse: 0.0533, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 01:02:40] (step=0000836) Train Loss mse: 0.0256, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 01:02:53] (step=0000837) Train Loss mse: 0.0375, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 01:03:09] (step=0000838) Train Loss mse: 0.0386, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 01:03:22] (step=0000839) Train Loss mse: 0.0483, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 01:03:34] (step=0000840) Train Loss mse: 0.0459, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 01:03:50] (step=0000841) Train Loss mse: 0.0493, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 01:04:04] (step=0000842) Train Loss mse: 0.0548, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 01:04:20] (step=0000843) Train Loss mse: 0.0532, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 01:04:36] (step=0000844) Train Loss mse: 0.0271, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 01:04:48] (step=0000845) Train Loss mse: 0.0427, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 01:05:02] (step=0000846) Train Loss mse: 0.0425, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 01:05:18] (step=0000847) Train Loss mse: 0.0339, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 01:05:34] (step=0000848) Train Loss mse: 0.0311, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 01:05:50] (step=0000849) Train Loss mse: 0.0434, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 01:06:01] (step=0000850) Train Loss mse: 0.0379, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
+[2026-01-03 01:06:14] (step=0000851) Train Loss mse: 0.0447, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 01:06:26] (step=0000852) Train Loss mse: 0.0516, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 01:06:38] (step=0000853) Train Loss mse: 0.0351, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
 FullyShardedDataParallel(
   (_fsdp_wrapped_module): Bagel(
     (language_model): Qwen2ForCausalLM(
@@ -920,121 +1035,6 @@ connector._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
 vit_pos_embed._fsdp_wrapped_module._flat_param False
 Preparing Dataset vlm_gym_jigsaw_mse_loss_only/vlm_gym_jigsaw_train
 Preparing Dataset vlm_gym_jigsaw_mse_loss_only/vlm_gym_jigsaw_train
-[2026-01-03 00:40:37] (step=0000739) Train Loss mse: 0.0488, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2026-01-03 00:40:50] (step=0000740) Train Loss mse: 0.0456, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 00:41:04] (step=0000741) Train Loss mse: 0.0518, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 00:41:16] (step=0000742) Train Loss mse: 0.0441, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 00:41:27] (step=0000743) Train Loss mse: 0.0542, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
-[2026-01-03 00:41:41] (step=0000744) Train Loss mse: 0.0310, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 00:41:53] (step=0000745) Train Loss mse: 0.0629, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 00:42:09] (step=0000746) Train Loss mse: 0.0451, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 00:42:23] (step=0000747) Train Loss mse: 0.0331, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 00:42:36] (step=0000748) Train Loss mse: 0.0387, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2026-01-03 00:42:47] (step=0000749) Train Loss mse: 0.0353, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
-[2026-01-03 00:42:59] (step=0000750) Train Loss mse: 0.0455, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
-[2026-01-03 00:43:12] (step=0000751) Train Loss mse: 0.0436, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 00:43:25] (step=0000752) Train Loss mse: 0.0441, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2026-01-03 00:43:41] (step=0000753) Train Loss mse: 0.0404, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 00:43:51] (step=0000754) Train Loss mse: 0.0379, Train Loss ce: 0.0000, Train Steps/Sec: 0.10,
-[2026-01-03 00:44:04] (step=0000755) Train Loss mse: 0.0340, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2026-01-03 00:44:20] (step=0000756) Train Loss mse: 0.0289, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 00:44:36] (step=0000757) Train Loss mse: 0.0436, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 00:44:52] (step=0000758) Train Loss mse: 0.0460, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 00:45:05] (step=0000759) Train Loss mse: 0.0317, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 00:45:22] (step=0000760) Train Loss mse: 0.0320, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 00:45:35] (step=0000761) Train Loss mse: 0.0383, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2026-01-03 00:45:46] (step=0000762) Train Loss mse: 0.0434, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
-[2026-01-03 00:46:03] (step=0000763) Train Loss mse: 0.0465, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 00:46:14] (step=0000764) Train Loss mse: 0.0320, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
-[2026-01-03 00:46:27] (step=0000765) Train Loss mse: 0.0423, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 00:46:40] (step=0000766) Train Loss mse: 0.0428, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2026-01-03 00:46:52] (step=0000767) Train Loss mse: 0.0392, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 00:47:06] (step=0000768) Train Loss mse: 0.0506, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 00:47:22] (step=0000769) Train Loss mse: 0.0424, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 00:47:34] (step=0000770) Train Loss mse: 0.0405, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 00:47:49] (step=0000771) Train Loss mse: 0.0442, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2026-01-03 00:48:02] (step=0000772) Train Loss mse: 0.0469, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2026-01-03 00:48:15] (step=0000773) Train Loss mse: 0.0456, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 00:48:28] (step=0000774) Train Loss mse: 0.0431, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 00:48:44] (step=0000775) Train Loss mse: 0.0369, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 00:49:00] (step=0000776) Train Loss mse: 0.0357, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 00:49:14] (step=0000777) Train Loss mse: 0.0466, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 00:49:30] (step=0000778) Train Loss mse: 0.0367, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 00:49:46] (step=0000779) Train Loss mse: 0.0452, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 00:49:56] (step=0000780) Train Loss mse: 0.0635, Train Loss ce: 0.0000, Train Steps/Sec: 0.10,
-[2026-01-03 00:50:08] (step=0000781) Train Loss mse: 0.0540, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 00:50:24] (step=0000782) Train Loss mse: 0.0468, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 00:50:36] (step=0000783) Train Loss mse: 0.0439, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 00:50:48] (step=0000784) Train Loss mse: 0.0506, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 00:51:00] (step=0000785) Train Loss mse: 0.0376, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 00:51:16] (step=0000786) Train Loss mse: 0.0326, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 00:51:29] (step=0000787) Train Loss mse: 0.0364, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 00:51:41] (step=0000788) Train Loss mse: 0.0292, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 00:51:57] (step=0000789) Train Loss mse: 0.0470, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 00:52:13] (step=0000790) Train Loss mse: 0.0311, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 00:52:27] (step=0000791) Train Loss mse: 0.0399, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 00:52:43] (step=0000792) Train Loss mse: 0.0250, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 00:52:55] (step=0000793) Train Loss mse: 0.0513, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 00:53:06] (step=0000794) Train Loss mse: 0.0426, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
-[2026-01-03 00:53:20] (step=0000795) Train Loss mse: 0.0637, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 00:53:33] (step=0000796) Train Loss mse: 0.0504, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2026-01-03 00:53:48] (step=0000797) Train Loss mse: 0.0467, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2026-01-03 00:54:00] (step=0000798) Train Loss mse: 0.0493, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
-[2026-01-03 00:54:16] (step=0000799) Train Loss mse: 0.0400, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 00:54:32] (step=0000800) Train Loss mse: 0.0519, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 00:54:44] (step=0000801) Train Loss mse: 0.0510, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 00:54:58] (step=0000802) Train Loss mse: 0.0425, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2026-01-03 00:55:09] (step=0000803) Train Loss mse: 0.0396, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
-[2026-01-03 00:55:19] (step=0000804) Train Loss mse: 0.0541, Train Loss ce: 0.0000, Train Steps/Sec: 0.11,
-[2026-01-03 00:55:32] (step=0000805) Train Loss mse: 0.0452, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 00:55:43] (step=0000806) Train Loss mse: 0.0527, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
-[2026-01-03 00:55:56] (step=0000807) Train Loss mse: 0.0441, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 00:56:13] (step=0000808) Train Loss mse: 0.0269, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 00:56:26] (step=0000809) Train Loss mse: 0.0419, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2026-01-03 00:56:39] (step=0000810) Train Loss mse: 0.0485, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 00:56:53] (step=0000811) Train Loss mse: 0.0351, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 00:57:09] (step=0000812) Train Loss mse: 0.0273, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 00:57:25] (step=0000813) Train Loss mse: 0.0412, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 00:57:37] (step=0000814) Train Loss mse: 0.0527, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 00:57:53] (step=0000815) Train Loss mse: 0.0316, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 00:58:06] (step=0000816) Train Loss mse: 0.0294, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2026-01-03 00:58:23] (step=0000817) Train Loss mse: 0.0337, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 00:58:39] (step=0000818) Train Loss mse: 0.0415, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 00:58:52] (step=0000819) Train Loss mse: 0.0356, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 00:59:05] (step=0000820) Train Loss mse: 0.0428, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2026-01-03 00:59:19] (step=0000821) Train Loss mse: 0.0456, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 00:59:33] (step=0000822) Train Loss mse: 0.0497, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2026-01-03 00:59:46] (step=0000823) Train Loss mse: 0.0428, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 01:00:02] (step=0000824) Train Loss mse: 0.0385, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 01:00:13] (step=0000825) Train Loss mse: 0.0499, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
-[2026-01-03 01:00:25] (step=0000826) Train Loss mse: 0.0360, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 01:00:39] (step=0000827) Train Loss mse: 0.0415, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2026-01-03 01:00:52] (step=0000828) Train Loss mse: 0.0479, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 01:01:05] (step=0000829) Train Loss mse: 0.0361, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 01:01:17] (step=0000830) Train Loss mse: 0.0509, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
-[2026-01-03 01:01:30] (step=0000831) Train Loss mse: 0.0425, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 01:01:43] (step=0000832) Train Loss mse: 0.0365, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 01:01:59] (step=0000833) Train Loss mse: 0.0427, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 01:02:11] (step=0000834) Train Loss mse: 0.0505, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 01:02:24] (step=0000835) Train Loss mse: 0.0533, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 01:02:40] (step=0000836) Train Loss mse: 0.0256, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 01:02:53] (step=0000837) Train Loss mse: 0.0375, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2026-01-03 01:03:09] (step=0000838) Train Loss mse: 0.0386, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 01:03:22] (step=0000839) Train Loss mse: 0.0483, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 01:03:34] (step=0000840) Train Loss mse: 0.0459, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 01:03:50] (step=0000841) Train Loss mse: 0.0493, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 01:04:04] (step=0000842) Train Loss mse: 0.0548, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2026-01-03 01:04:20] (step=0000843) Train Loss mse: 0.0532, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 01:04:36] (step=0000844) Train Loss mse: 0.0271, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 01:04:48] (step=0000845) Train Loss mse: 0.0427, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 01:05:02] (step=0000846) Train Loss mse: 0.0425, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2026-01-03 01:05:18] (step=0000847) Train Loss mse: 0.0339, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 01:05:34] (step=0000848) Train Loss mse: 0.0311, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 01:05:50] (step=0000849) Train Loss mse: 0.0434, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 01:06:01] (step=0000850) Train Loss mse: 0.0379, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
-[2026-01-03 01:06:14] (step=0000851) Train Loss mse: 0.0447, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 01:06:26] (step=0000852) Train Loss mse: 0.0516, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 01:06:38] (step=0000853) Train Loss mse: 0.0351, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
 [2026-01-03 01:06:51] (step=0000854) Train Loss mse: 0.0607, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
 [2026-01-03 01:07:02] (step=0000855) Train Loss mse: 0.0458, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
 [2026-01-03 01:07:15] (step=0000856) Train Loss mse: 0.0375, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
@@ -3098,7 +3098,86 @@ Preparing Dataset vlm_gym_jigsaw_mse_loss_only/vlm_gym_jigsaw_train
 [2026-01-03 09:06:11] (step=0002911) Train Loss mse: 0.0372, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
 [2026-01-03 09:06:24] (step=0002912) Train Loss mse: 0.0293, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
 [2026-01-03 09:06:37] (step=0002913) Train Loss mse: 0.0295, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 09:06:54
+[2026-01-03 09:06:54] (step=0002914) Train Loss mse: 0.0248, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 09:07:07] (step=0002915) Train Loss mse: 0.0321, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 09:07:23] (step=0002916) Train Loss mse: 0.0410, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 09:07:34] (step=0002917) Train Loss mse: 0.0278, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
+[2026-01-03 09:07:48] (step=0002918) Train Loss mse: 0.0212, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 09:08:02] (step=0002919) Train Loss mse: 0.0319, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 09:08:15] (step=0002920) Train Loss mse: 0.0301, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 09:08:28] (step=0002921) Train Loss mse: 0.0250, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 09:08:42] (step=0002922) Train Loss mse: 0.0364, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 09:08:54] (step=0002923) Train Loss mse: 0.0326, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
+[2026-01-03 09:09:07] (step=0002924) Train Loss mse: 0.0247, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 09:09:20] (step=0002925) Train Loss mse: 0.0255, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 09:09:33] (step=0002926) Train Loss mse: 0.0297, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 09:09:43] (step=0002927) Train Loss mse: 0.0442, Train Loss ce: 0.0000, Train Steps/Sec: 0.10,
+[2026-01-03 09:09:56] (step=0002928) Train Loss mse: 0.0392, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 09:10:09] (step=0002929) Train Loss mse: 0.0275, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 09:10:22] (step=0002930) Train Loss mse: 0.0327, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 09:10:34] (step=0002931) Train Loss mse: 0.0247, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 09:10:47] (step=0002932) Train Loss mse: 0.0272, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 09:11:00] (step=0002933) Train Loss mse: 0.0340, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 09:11:13] (step=0002934) Train Loss mse: 0.0292, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 09:11:30] (step=0002935) Train Loss mse: 0.0287, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 09:11:46] (step=0002936) Train Loss mse: 0.0377, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 09:12:03] (step=0002937) Train Loss mse: 0.0310, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 09:12:16] (step=0002938) Train Loss mse: 0.0239, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 09:12:28] (step=0002939) Train Loss mse: 0.0395, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
+[2026-01-03 09:12:44] (step=0002940) Train Loss mse: 0.0276, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 09:12:57] (step=0002941) Train Loss mse: 0.0327, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 09:13:10] (step=0002942) Train Loss mse: 0.0221, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 09:13:26] (step=0002943) Train Loss mse: 0.0456, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 09:13:41] (step=0002944) Train Loss mse: 0.0258, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 09:13:55] (step=0002945) Train Loss mse: 0.0255, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 09:14:07] (step=0002946) Train Loss mse: 0.0385, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 09:14:21] (step=0002947) Train Loss mse: 0.0289, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 09:14:35] (step=0002948) Train Loss mse: 0.0198, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 09:14:50] (step=0002949) Train Loss mse: 0.0257, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 09:15:06] (step=0002950) Train Loss mse: 0.0224, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 09:15:16] (step=0002951) Train Loss mse: 0.0208, Train Loss ce: 0.0000, Train Steps/Sec: 0.11,
+[2026-01-03 09:15:29] (step=0002952) Train Loss mse: 0.0373, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 09:15:41] (step=0002953) Train Loss mse: 0.0369, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 09:15:55] (step=0002954) Train Loss mse: 0.0235, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 09:16:11] (step=0002955) Train Loss mse: 0.0317, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 09:16:24] (step=0002956) Train Loss mse: 0.0285, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 09:16:40] (step=0002957) Train Loss mse: 0.0308, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 09:16:54] (step=0002958) Train Loss mse: 0.0306, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 09:17:05] (step=0002959) Train Loss mse: 0.0333, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
+[2026-01-03 09:17:22] (step=0002960) Train Loss mse: 0.0246, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 09:17:35] (step=0002961) Train Loss mse: 0.0196, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 09:17:49] (step=0002962) Train Loss mse: 0.0269, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 09:18:02] (step=0002963) Train Loss mse: 0.0276, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 09:18:18] (step=0002964) Train Loss mse: 0.0235, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 09:18:34] (step=0002965) Train Loss mse: 0.0461, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 09:18:45] (step=0002966) Train Loss mse: 0.0278, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
+[2026-01-03 09:18:58] (step=0002967) Train Loss mse: 0.0296, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 09:19:12] (step=0002968) Train Loss mse: 0.0292, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 09:19:28] (step=0002969) Train Loss mse: 0.0320, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 09:19:44] (step=0002970) Train Loss mse: 0.0255, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 09:19:57] (step=0002971) Train Loss mse: 0.0284, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 09:20:10] (step=0002972) Train Loss mse: 0.0411, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 09:20:26] (step=0002973) Train Loss mse: 0.0230, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 09:20:43] (step=0002974) Train Loss mse: 0.0307, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 09:20:56] (step=0002975) Train Loss mse: 0.0432, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 09:21:12] (step=0002976) Train Loss mse: 0.0236, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 09:21:23] (step=0002977) Train Loss mse: 0.0258, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
+[2026-01-03 09:21:40] (step=0002978) Train Loss mse: 0.0384, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 09:21:53] (step=0002979) Train Loss mse: 0.0200, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 09:22:06] (step=0002980) Train Loss mse: 0.0263, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 09:22:18] (step=0002981) Train Loss mse: 0.0279, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
+[2026-01-03 09:22:34] (step=0002982) Train Loss mse: 0.0299, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 09:22:50] (step=0002983) Train Loss mse: 0.0211, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 09:23:03] (step=0002984) Train Loss mse: 0.0263, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 09:23:17] (step=0002985) Train Loss mse: 0.0286, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 09:23:30] (step=0002986) Train Loss mse: 0.0498, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 09:23:43] (step=0002987) Train Loss mse: 0.0259, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 09:23:56] (step=0002988) Train Loss mse: 0.0292, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 09:24:09] (step=0002989) Train Loss mse: 0.0224, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 09:24:25] (step=0002990) Train Loss mse: 0.0231, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 09:24:41] (step=0002991) Train Loss mse: 0.0408, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 09:24:57] (step=0002992) Train Loss mse: 0.0458, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 09:25:13] (step=0002993) Train Loss mse: 0.0287, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
 [2026-01-03 09:25:26] (step=0002994) Train Loss mse: 0.0343, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
 [2026-01-03 09:25:38] (step=0002995) Train Loss mse: 0.0318, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
 [2026-01-03 09:25:51] (step=0002996) Train Loss mse: 0.0319, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
@@ -4109,67 +4188,76 @@ Preparing Dataset vlm_gym_jigsaw_mse_loss_only/vlm_gym_jigsaw_train
 [2026-01-03 13:19:24] (step=0004001) Train Loss mse: 0.0321, Train Loss ce: 0.0000, Train Steps/Sec: 0.01,
 [2026-01-03 13:19:40] (step=0004002) Train Loss mse: 0.0269, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
 [2026-01-03 13:19:53] (step=0004003) Train Loss mse: 0.0293, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 09:10:47] (step=0002932) Train Loss mse: 0.0272, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 09:11:00] (step=0002933) Train Loss mse: 0.0340, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 09:11:13] (step=0002934) Train Loss mse: 0.0292, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 09:11:30] (step=0002935) Train Loss mse: 0.0287, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 09:11:46] (step=0002936) Train Loss mse: 0.0377, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 09:12:03] (step=0002937) Train Loss mse: 0.0310, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 09:12:16] (step=0002938) Train Loss mse: 0.0239, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 09:12:28] (step=0002939) Train Loss mse: 0.0395, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
-[2026-01-03 09:12:44] (step=0002940) Train Loss mse: 0.0276, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 09:12:57] (step=0002941) Train Loss mse: 0.0327, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 09:13:10] (step=0002942) Train Loss mse: 0.0221, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 09:13:26] (step=0002943) Train Loss mse: 0.0456, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 09:13:41] (step=0002944) Train Loss mse: 0.0258, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2026-01-03 09:13:55] (step=0002945) Train Loss mse: 0.0255, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2026-01-03 09:14:07] (step=0002946) Train Loss mse: 0.0385, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 09:14:21] (step=0002947) Train Loss mse: 0.0289, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2026-01-03 09:14:35] (step=0002948) Train Loss mse: 0.0198, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2026-01-03 09:14:50] (step=0002949) Train Loss mse: 0.0257, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2026-01-03 09:15:06] (step=0002950) Train Loss mse: 0.0224, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 09:15:16] (step=0002951) Train Loss mse: 0.0208, Train Loss ce: 0.0000, Train Steps/Sec: 0.11,
-[2026-01-03 09:15:29] (step=0002952) Train Loss mse: 0.0373, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 09:15:41] (step=0002953) Train Loss mse: 0.0369, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 09:15:55] (step=0002954) Train Loss mse: 0.0235, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 09:16:11] (step=0002955) Train Loss mse: 0.0317, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 09:16:24] (step=0002956) Train Loss mse: 0.0285, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2026-01-03 09:16:40] (step=0002957) Train Loss mse: 0.0308, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 09:16:54] (step=0002958) Train Loss mse: 0.0306, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2026-01-03 09:17:05] (step=0002959) Train Loss mse: 0.0333, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
-[2026-01-03 09:17:22] (step=0002960) Train Loss mse: 0.0246, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 09:17:35] (step=0002961) Train Loss mse: 0.0196, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2026-01-03 09:17:49] (step=0002962) Train Loss mse: 0.0269, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2026-01-03 09:18:02] (step=0002963) Train Loss mse: 0.0276, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 09:18:18] (step=0002964) Train Loss mse: 0.0235, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 09:18:34] (step=0002965) Train Loss mse: 0.0461, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 09:18:45] (step=0002966) Train Loss mse: 0.0278, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
-[2026-01-03 09:18:58] (step=0002967) Train Loss mse: 0.0296, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 09:19:12] (step=0002968) Train Loss mse: 0.0292, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2026-01-03 09:19:28] (step=0002969) Train Loss mse: 0.0320, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 09:19:44] (step=0002970) Train Loss mse: 0.0255, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 09:19:57] (step=0002971) Train Loss mse: 0.0284, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 09:20:10] (step=0002972) Train Loss mse: 0.0411, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 09:20:26] (step=0002973) Train Loss mse: 0.0230, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 09:20:43] (step=0002974) Train Loss mse: 0.0307, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 09:20:56] (step=0002975) Train Loss mse: 0.0432, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2026-01-03 09:21:12] (step=0002976) Train Loss mse: 0.0236, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 09:21:23] (step=0002977) Train Loss mse: 0.0258, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
-[2026-01-03 09:21:40] (step=0002978) Train Loss mse: 0.0384, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 09:21:53] (step=0002979) Train Loss mse: 0.0200, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 09:22:06] (step=0002980) Train Loss mse: 0.0263, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2026-01-03 09:22:18] (step=0002981) Train Loss mse: 0.0279, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
-[2026-01-03 09:22:34] (step=0002982) Train Loss mse: 0.0299, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 09:22:50] (step=0002983) Train Loss mse: 0.0211, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 09:23:03] (step=0002984) Train Loss mse: 0.0263, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 09:23:17] (step=0002985) Train Loss mse: 0.0286, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 09:23:30] (step=0002986) Train Loss mse: 0.0498, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 09:23:43] (step=0002987) Train Loss mse: 0.0259, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 09:23:56] (step=0002988) Train Loss mse: 0.0292, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
-[2026-01-03 09:24:09] (step=0002989) Train Loss mse: 0.0224, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
-[2026-01-03 09:24:25] (step=0002990) Train Loss mse: 0.0231, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 09:24:41] (step=0002991) Train Loss mse: 0.0408, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
-[2026-01-03 09:24:57] (step=0002992) Train Loss mse: 0.0458, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 13:20:06] (step=0004004) Train Loss mse: 0.0326, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 13:20:22] (step=0004005) Train Loss mse: 0.0287, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 13:20:35] (step=0004006) Train Loss mse: 0.0288, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 13:20:52] (step=0004007) Train Loss mse: 0.0311, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 13:21:05] (step=0004008) Train Loss mse: 0.0316, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 13:21:17] (step=0004009) Train Loss mse: 0.0310, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 13:21:30] (step=0004010) Train Loss mse: 0.0216, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 13:21:44] (step=0004011) Train Loss mse: 0.0219, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 13:21:57] (step=0004012) Train Loss mse: 0.0231, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 13:22:08] (step=0004013) Train Loss mse: 0.0348, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
+[2026-01-03 13:22:20] (step=0004014) Train Loss mse: 0.0390, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
+[2026-01-03 13:22:33] (step=0004015) Train Loss mse: 0.0276, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 13:22:46] (step=0004016) Train Loss mse: 0.0321, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 13:22:58] (step=0004017) Train Loss mse: 0.0268, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 13:23:11] (step=0004018) Train Loss mse: 0.0239, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 13:23:24] (step=0004019) Train Loss mse: 0.0309, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 13:23:41] (step=0004020) Train Loss mse: 0.0306, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 13:23:57] (step=0004021) Train Loss mse: 0.0226, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 13:24:10] (step=0004022) Train Loss mse: 0.0334, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 13:24:22] (step=0004023) Train Loss mse: 0.0287, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 13:24:35] (step=0004024) Train Loss mse: 0.0285, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 13:24:51] (step=0004025) Train Loss mse: 0.0191, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 13:25:07] (step=0004026) Train Loss mse: 0.0229, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 13:25:20] (step=0004027) Train Loss mse: 0.0278, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 13:25:36] (step=0004028) Train Loss mse: 0.0277, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 13:25:52] (step=0004029) Train Loss mse: 0.0230, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 13:26:08] (step=0004030) Train Loss mse: 0.0207, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 13:26:24] (step=0004031) Train Loss mse: 0.0251, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 13:26:35] (step=0004032) Train Loss mse: 0.0274, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
+[2026-01-03 13:26:49] (step=0004033) Train Loss mse: 0.0324, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 13:27:02] (step=0004034) Train Loss mse: 0.0267, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 13:27:14] (step=0004035) Train Loss mse: 0.0361, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 13:27:27] (step=0004036) Train Loss mse: 0.0190, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 13:27:42] (step=0004037) Train Loss mse: 0.0312, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 13:27:58] (step=0004038) Train Loss mse: 0.0341, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 13:28:11] (step=0004039) Train Loss mse: 0.0216, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 13:28:27] (step=0004040) Train Loss mse: 0.0233, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 13:28:40] (step=0004041) Train Loss mse: 0.0216, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 13:28:56] (step=0004042) Train Loss mse: 0.0303, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 13:29:09] (step=0004043) Train Loss mse: 0.0208, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 13:29:25] (step=0004044) Train Loss mse: 0.0203, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 13:29:41] (step=0004045) Train Loss mse: 0.0265, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 13:29:57] (step=0004046) Train Loss mse: 0.0281, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 13:30:10] (step=0004047) Train Loss mse: 0.0387, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 13:30:23] (step=0004048) Train Loss mse: 0.0330, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 13:30:37] (step=0004049) Train Loss mse: 0.0340, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 13:30:50] (step=0004050) Train Loss mse: 0.0228, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 13:31:01] (step=0004051) Train Loss mse: 0.0361, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
+[2026-01-03 13:31:17] (step=0004052) Train Loss mse: 0.0259, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 13:31:30] (step=0004053) Train Loss mse: 0.0337, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 13:31:43] (step=0004054) Train Loss mse: 0.0250, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 13:31:57] (step=0004055) Train Loss mse: 0.0260, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 13:32:08] (step=0004056) Train Loss mse: 0.0335, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
+[2026-01-03 13:32:20] (step=0004057) Train Loss mse: 0.0278, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 13:32:33] (step=0004058) Train Loss mse: 0.0328, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 13:32:49] (step=0004059) Train Loss mse: 0.0268, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 13:33:05] (step=0004060) Train Loss mse: 0.0228, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 13:33:19] (step=0004061) Train Loss mse: 0.0318, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 13:33:32] (step=0004062) Train Loss mse: 0.0223, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 13:33:48] (step=0004063) Train Loss mse: 0.0239, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
+[2026-01-03 13:34:03] (step=0004064) Train Loss mse: 0.0381, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 13:34:14] (step=0004065) Train Loss mse: 0.0237, Train Loss ce: 0.0000, Train Steps/Sec: 0.09,
+[2026-01-03 13:34:28] (step=0004066) Train Loss mse: 0.0223, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 13:34:41] (step=0004067) Train Loss mse: 0.0347, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 13:34:54] (step=0004068) Train Loss mse: 0.0303, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 13:35:08] (step=0004069) Train Loss mse: 0.0304, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 13:35:21] (step=0004070) Train Loss mse: 0.0254, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 13:35:35] (step=0004071) Train Loss mse: 0.0207, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 13:35:47] (step=0004072) Train Loss mse: 0.0320, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
+[2026-01-03 13:35:59] (step=0004073) Train Loss mse: 0.0377, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
 [2026-01-03 13:36:13] (step=0004074) Train Loss mse: 0.0289, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
 [2026-01-03 13:36:25] (step=0004075) Train Loss mse: 0.0329, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
 [2026-01-03 13:36:38] (step=0004076) Train Loss mse: 0.0261, Train Loss ce: 0.0000, Train Steps/Sec: 0.08,
@@ -5098,5 +5186,4 @@ Preparing Dataset vlm_gym_jigsaw_mse_loss_only/vlm_gym_jigsaw_train
 [2026-01-03 17:08:32] (step=0004999) Train Loss mse: 0.0250, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
 [2026-01-03 17:08:48] (step=0005000) Train Loss mse: 0.0230, Train Loss ce: 0.0000, Train Steps/Sec: 0.06,
 [2026-01-03 17:10:07] Saving checkpoint to /dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_mse_only/0005000.
-[2026-01-03 17:12:43] Done!
-[2026-01-03 13:36:13] (step=0004074) Train Loss mse: 0.0289, Train Loss ce: 0.0000, Train Steps/Sec: 0.07,
+[2026-01-03 17:12:43] Done!

checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_mse_only/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_mse_only/wandb/offline-run-20260104_093200-vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0/files/requirements.txt

@@ -0,0 +1,354 @@
+Brotli==1.1.0
+MarkupSafe==3.0.2
+PySocks==1.7.1
+PyYAML==6.0.2
+archspec==0.2.3
+asttokens==2.4.1
+astunparse==1.6.3
+attrs==24.2.0
+beautifulsoup4==4.12.3
+boltons==24.0.0
+certifi==2024.8.30
+chardet==5.2.0
+charset-normalizer==3.4.0
+click==8.1.7
+colorama==0.4.6
+conda==24.9.2
+conda-build==24.9.0
+conda_index==0.5.0
+conda-libmamba-solver==24.9.0
+conda-package-handling==2.4.0
+conda_package_streaming==0.11.0
+decorator==5.1.1
+distro==1.9.0
+dnspython==2.7.0
+exceptiongroup==1.2.2
+executing==2.1.0
+expecttest==0.2.1
+filelock==3.16.1
+frozendict==2.4.6
+fsspec==2024.10.0
+h2==4.1.0
+hpack==4.0.0
+hyperframe==6.0.1
+hypothesis==6.115.5
+idna==3.10
+importlib_resources==6.4.5
+ipython==8.29.0
+jedi==0.19.1
+Jinja2==3.1.4
+jsonpatch==1.33
+jsonpointer==3.0.0
+jsonschema==4.23.0
+jsonschema-specifications==2024.10.1
+libarchive-c==5.1
+libmambapy==1.5.10
+lief==0.14.1
+lintrunner==0.12.5
+mamba==1.5.10
+matplotlib-inline==0.1.7
+menuinst==2.1.2
+more-itertools==10.5.0
+mpmath==1.3.0
+networkx==3.4.2
+ninja==1.11.1.1
+nvidia-cublas-cu12==12.4.5.8
+nvidia-cuda-cupti-cu12==12.4.127
+nvidia-cuda-nvrtc-cu12==12.4.127
+nvidia-cuda-runtime-cu12==12.4.127
+nvidia-cudnn-cu12==9.1.0.70
+nvidia-cufft-cu12==11.2.1.3
+nvidia-curand-cu12==10.3.5.147
+nvidia-cusolver-cu12==11.6.1.9
+nvidia-cusparse-cu12==12.3.1.170
+nvidia-nccl-cu12==2.21.5
+nvidia-nvjitlink-cu12==12.4.127
+nvidia-nvtx-cu12==12.4.127
+optree==0.13.0
+parso==0.8.4
+pexpect==4.9.0
+pickleshare==0.7.5
+pillow==10.2.0
+pkginfo==1.11.2
+pkgutil_resolve_name==1.3.10
+platformdirs==4.3.6
+pluggy==1.5.0
+prompt_toolkit==3.0.48
+psutil==6.1.0
+ptyprocess==0.7.0
+pure_eval==0.2.3
+pycosat==0.6.6
+pycparser==2.22
+Pygments==2.18.0
+python-etcd==0.4.5
+pytz==2024.2
+referencing==0.35.1
+requests==2.32.3
+rpds-py==0.20.0
+ruamel.yaml==0.18.6
+ruamel.yaml.clib==0.2.8
+six==1.16.0
+sortedcontainers==2.4.0
+soupsieve==2.5
+stack-data==0.6.2
+sympy==1.13.1
+torchaudio==2.5.1+cu124
+torchelastic==0.2.2
+tqdm==4.66.5
+traitlets==5.14.3
+triton==3.1.0
+truststore==0.9.2
+types-dataclasses==0.6.6
+urllib3==2.2.3
+wcwidth==0.2.13
+zipp==3.20.2
+zstandard==0.23.0
+numpy==1.24.4
+imgcat==0.6.0
+decord==0.6.0
+flash_attn==2.5.8
+contourpy==1.3.2
+cycler==0.12.1
+fonttools==4.61.1
+huggingface-hub==0.29.1
+kiwisolver==1.4.9
+matplotlib==3.7.0
+opencv-python==4.7.0.72
+pyarrow==11.0.0
+pyparsing==3.2.5
+safetensors==0.4.5
+scipy==1.10.1
+sentencepiece==0.1.99
+torch==2.5.1
+torchvision==0.20.1
+transformers==4.49.0
+pip==25.3
+setuptools==80.9.0
+wheel==0.45.1
+Pebble==5.1.3
+accelerate==1.12.0
+addftool==0.2.13
+aiohappyeyeballs==2.6.1
+aiohttp==3.13.2
+aiohttp-cors==0.8.1
+aiosignal==1.4.0
+airportsdata==20250909
+annotated-doc==0.0.4
+annotated-types==0.7.0
+antlr4-python3-runtime==4.9.3
+bcrypt==5.0.0
+blobfile==3.0.0
+cffi==2.0.0
+cloudpickle==3.1.2
+codetiming==1.4.0
+colorful==0.5.8
+compressed-tensors==0.12.2
+cryptography==46.0.3
+cuda-bindings==13.1.1
+cuda-pathfinder==1.3.3
+cuda-python==13.1.1
+datasets==4.4.1
+Deprecated==1.3.1
+diskcache==5.6.3
+distlib==0.4.0
+docstring_parser==0.17.0
+easydict==1.13
+fabric==3.2.2
+fastapi==0.124.4
+fire==0.7.1
+flashinfer-python==0.2.5
+frozenlist==1.8.0
+gevent==25.9.1
+gitdb==4.0.12
+GitPython==3.1.45
+google-api-core==2.28.1
+google-auth==2.43.0
+google-cloud-aiplatform==1.130.0
+google-cloud-bigquery==3.38.0
+google-cloud-core==2.5.0
+google-cloud-resource-manager==1.15.0
+google-cloud-storage==3.7.0
+google-crc32c==1.7.1
+google-genai==1.55.0
+google-resumable-media==2.8.0
+googleapis-common-protos==1.72.0
+greenlet==3.3.0
+grpc-google-iam-v1==0.14.3
+grpcio==1.76.0
+grpcio-status==1.76.0
+hf_transfer==0.1.9
+hf-xet==1.2.0
+hydra-core==1.3.2
+importlib_metadata==8.7.0
+interegular==0.3.3
+invoke==2.2.1
+jiter==0.12.0
+joblib==1.5.2
+jsonlines==4.0.0
+lark==1.3.1
+latex2sympy2==1.5.4
+latex2sympy2_extended==1.10.2
+libtmux==0.52.1
+llguidance==0.7.30
+loguru==0.7.3
+lxml==6.0.2
+math-verify==0.8.0
+modelscope==1.33.0
+msgpack==1.1.2
+msgspec==0.20.0
+multidict==6.7.0
+multiprocess==0.70.18
+nvidia-cusparselt-cu12==0.6.2
+nvidia-ml-py==13.590.44
+omegaconf==2.3.0
+openai==2.11.0
+opencensus==0.11.4
+opencensus-context==0.1.3
+opentelemetry-api==1.39.1
+opentelemetry-exporter-prometheus==0.60b1
+opentelemetry-proto==1.39.1
+opentelemetry-sdk==1.39.1
+opentelemetry-semantic-conventions==0.60b1
+orjson==3.11.5
+outlines==0.1.11
+outlines_core==0.1.26
+packaging==25.0
+pandas==2.3.3
+parallel-ssh==2.16.0.post1
+paramiko==4.0.0
+partial-json-parser==0.2.1.1.post7
+peft==0.18.0
+propcache==0.4.1
+proto-plus==1.26.1
+protobuf==6.33.2
+py-spy==0.4.1
+pyasn1==0.6.1
+pyasn1_modules==0.4.2
+pybind11==3.0.1
+pycountry==24.6.1
+pycryptodomex==3.23.0
+pydantic==2.12.5
+pydantic_core==2.41.5
+pylatexenc==2.10
+PyNaCl==1.6.1
+pynvml==13.0.1
+python-multipart==0.0.20
+ray==2.52.1
+regex==2025.11.3
+rsa==4.9.1
+scikit-learn==1.8.0
+sentence-transformers==5.2.0
+sentry-sdk==2.47.0
+setproctitle==1.3.7
+sgl-kernel==0.1.4
+sglang==0.4.6.post5
+shapely==2.1.2
+smart_open==7.5.0
+smmap==5.0.2
+sniffio==1.3.1
+soundfile==0.13.1
+ssh2-python==1.2.0.post1
+ssh-python==1.2.0.post1
+starlette==0.50.0
+tabulate==0.9.0
+tenacity==9.1.2
+tensorboardX==2.6.4
+tensordict==0.6.2
+termcolor==3.2.0
+threadpoolctl==3.6.0
+tiktoken==0.12.0
+timeout-decorator==0.5.0
+tmuxp==1.61.0
+tokenizers==0.21.4
+torch_memory_saver==0.0.9
+torchao==0.9.0
+torchdata==0.11.0
+typing-inspection==0.4.2
+uvicorn==0.38.0
+uvloop==0.22.1
+virtualenv==20.35.4
+wandb==0.23.1
+websockets==15.0.1
+word2number==1.1
+wrapt==2.0.1
+xgrammar==0.1.19
+xxhash==3.6.0
+yarl==1.22.0
+zope.event==6.1
+zope.interface==8.1.1
+cachetools==6.2.3
+dill==0.4.0
+inflect==7.5.0
+lazy_loader==0.4
+rp==0.1.1333
+stackprinter==0.2.12
+typeguard==4.4.4
+typing_extensions==4.15.0
+asciinema==2.4.0
+einops==0.8.1
+Send2Trash==1.8.3
+anyio==4.12.0
+argon2-cffi==25.1.0
+argon2-cffi-bindings==25.1.0
+arrow==1.4.0
+async-lru==2.0.5
+babel==2.17.0
+bleach==6.3.0
+comm==0.2.3
+debugpy==1.8.18
+defusedxml==0.7.1
+fastjsonschema==2.21.2
+fqdn==1.5.1
+h11==0.16.0
+httpcore==1.0.9
+httpx==0.28.1
+ipykernel==7.1.0
+isoduration==20.11.0
+json5==0.12.1
+jupyter_client==8.7.0
+jupyter_core==5.9.1
+jupyter-events==0.12.0
+jupyter-lsp==2.3.0
+jupyter_server==2.17.0
+jupyter_server_terminals==0.5.3
+jupyterlab==4.5.0
+jupyterlab_pygments==0.3.0
+jupyterlab_server==2.28.0
+mistune==3.1.4
+nbclient==0.10.2
+nbconvert==7.16.6
+nbformat==5.10.4
+nest-asyncio==1.6.0
+notebook_shim==0.2.4
+overrides==7.7.0
+pandocfilters==1.5.1
+prometheus_client==0.23.1
+python-dateutil==2.9.0.post0
+python-json-logger==4.0.0
+pyzmq==27.1.0
+rfc3339-validator==0.1.4
+rfc3986-validator==0.1.1
+terminado==0.18.1
+tinycss2==1.4.0
+tornado==6.5.3
+tzdata==2025.3
+uri-template==1.3.0
+webcolors==25.10.0
+webencodings==0.5.1
+websocket-client==1.9.0
+autocommand==2.2.2
+backports.tarfile==1.2.0
+importlib_metadata==8.0.0
+inflect==7.3.1
+jaraco.collections==5.1.0
+jaraco.context==5.3.0
+jaraco.functools==4.0.1
+jaraco.text==3.12.1
+more-itertools==10.3.0
+packaging==24.2
+platformdirs==4.2.2
+tomli==2.0.1
+typeguard==4.3.0
+typing_extensions==4.12.2
+wheel==0.45.1
+zipp==3.19.2

checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_mse_only/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_mse_only/wandb/offline-run-20260104_093200-vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0/logs/debug-core.log

@@ -0,0 +1,7 @@
+{"time":"2026-01-04T09:32:00.835208263Z","level":"INFO","msg":"main: starting server","port-filename":"/tmp/tmpc0hvyn70/port-13221.txt","pid":13221,"log-level":0,"disable-analytics":false,"shutdown-on-parent-exit":false,"enable-dcgm-profiling":false}
+{"time":"2026-01-04T09:32:00.838086455Z","level":"INFO","msg":"server: will exit if parent process dies","ppid":13221}
+{"time":"2026-01-04T09:32:00.838143904Z","level":"INFO","msg":"server: accepting connections","addr":{"Name":"/tmp/wandb-13221-13288-2426694750/socket","Net":"unix"}}
+{"time":"2026-01-04T09:32:00.939441588Z","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"1(@)"}
+{"time":"2026-01-04T09:32:00.950369146Z","level":"INFO","msg":"handleInformInit: received","streamId":"vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0","id":"1(@)"}
+{"time":"2026-01-04T09:32:01.443736617Z","level":"INFO","msg":"handleInformInit: stream started","streamId":"vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0","id":"1(@)"}
+{"time":"2026-01-04T09:37:29.649420411Z","level":"INFO","msg":"server: parent process exited, terminating service process"}

checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_mse_only/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_mse_only/wandb/offline-run-20260104_093200-vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0/logs/debug-internal.log

@@ -0,0 +1,8 @@
+{"time":"2026-01-04T09:32:01.115692117Z","level":"INFO","msg":"stream: starting","core version":"0.23.1"}
+{"time":"2026-01-04T09:32:01.44318528Z","level":"WARN","msg":"featurechecker: GraphQL client is nil, skipping feature loading"}
+{"time":"2026-01-04T09:32:01.443259623Z","level":"INFO","msg":"stream: created new stream","id":"vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0"}
+{"time":"2026-01-04T09:32:01.443726951Z","level":"INFO","msg":"stream: started","id":"vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0"}
+{"time":"2026-01-04T09:32:01.44372382Z","level":"INFO","msg":"handler: started","stream_id":"vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0"}
+{"time":"2026-01-04T09:32:01.443767451Z","level":"INFO","msg":"writer: started","stream_id":"vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0"}
+{"time":"2026-01-04T09:32:01.443797476Z","level":"INFO","msg":"sender: started","stream_id":"vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0"}
+{"time":"2026-01-04T09:32:01.444535545Z","level":"WARN","msg":"runupserter: server does not expand metric globs but the x_server_side_expand_glob_metrics setting is set; ignoring"}

checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_mse_only/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_mse_only/wandb/offline-run-20260104_093200-vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0/logs/debug.log

@@ -0,0 +1,24 @@
+2026-01-04 09:32:00,479 INFO    MainThread:13221 [wandb_setup.py:_flush():80] Current SDK version is 0.23.1
+2026-01-04 09:32:00,479 INFO    MainThread:13221 [wandb_setup.py:_flush():80] Configure stats pid to 13221
+2026-01-04 09:32:00,479 INFO    MainThread:13221 [wandb_setup.py:_flush():80] Loading settings from /home/clouduser/.config/wandb/settings
+2026-01-04 09:32:00,479 INFO    MainThread:13221 [wandb_setup.py:_flush():80] Loading settings from /home/clouduser/Code/Github/unified_world_model/wandb/settings
+2026-01-04 09:32:00,479 INFO    MainThread:13221 [wandb_setup.py:_flush():80] Loading settings from environment variables
+2026-01-04 09:32:00,479 INFO    MainThread:13221 [wandb_init.py:setup_run_log_directory():714] Logging user logs to /dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_mse_only/wandb/offline-run-20260104_093200-vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0/logs/debug.log
+2026-01-04 09:32:00,479 INFO    MainThread:13221 [wandb_init.py:setup_run_log_directory():715] Logging internal logs to /dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_mse_only/wandb/offline-run-20260104_093200-vlm_gym_jigsaw_one_img_lr2e_5_mse_only-run0/logs/debug-internal.log
+2026-01-04 09:32:00,479 INFO    MainThread:13221 [wandb_init.py:init():841] calling init triggers
+2026-01-04 09:32:00,480 INFO    MainThread:13221 [wandb_init.py:init():846] wandb.init called with sweep_config: {}
+config: {'_wandb': {}}
+2026-01-04 09:32:00,480 INFO    MainThread:13221 [wandb_init.py:init():889] starting backend
+2026-01-04 09:32:00,939 INFO    MainThread:13221 [wandb_init.py:init():892] sending inform_init request
+2026-01-04 09:32:00,947 INFO    MainThread:13221 [wandb_init.py:init():900] backend started and connected
+2026-01-04 09:32:00,949 INFO    MainThread:13221 [wandb_init.py:init():970] updated telemetry
+2026-01-04 09:32:00,956 INFO    MainThread:13221 [wandb_init.py:init():994] communicating run to backend with 90.0 second timeout
+2026-01-04 09:32:01,446 INFO    MainThread:13221 [wandb_init.py:init():1041] starting run threads in backend
+2026-01-04 09:32:01,964 INFO    MainThread:13221 [wandb_run.py:_console_start():2521] atexit reg
+2026-01-04 09:32:01,964 INFO    MainThread:13221 [wandb_run.py:_redirect():2369] redirect: wrap_raw
+2026-01-04 09:32:01,964 INFO    MainThread:13221 [wandb_run.py:_redirect():2438] Wrapping output streams.
+2026-01-04 09:32:01,964 INFO    MainThread:13221 [wandb_run.py:_redirect():2461] Redirects installed.
+2026-01-04 09:32:01,968 INFO    MainThread:13221 [wandb_init.py:init():1081] run started, returning control to user process
+2026-01-04 09:32:01,969 INFO    MainThread:13221 [wandb_run.py:_config_callback():1396] config_cb None None {'visual_gen': True, 'visual_und': True, 'results_dir': 'results', 'checkpoint_dir': '/dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_mse_only', 'wandb_project': 'bagel', 'wandb_name': 'vlm_gym_jigsaw_one_img_lr2e_5_mse_only', 'wandb_runid': '0', 'wandb_resume': 'allow', 'wandb_offline': True, 'wandb_dir': '/dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_mse_only', 'global_seed': 4396, 'auto_resume': False, 'resume_from': '/home/clouduser/Code/Models/BAGEL-7B-MoT', 'resume_model_only': True, 'finetune_from_ema': True, 'finetune_from_hf': True, 'log_every': 1, 'save_every': 2500, 'total_steps': 5000, 'warmup_steps': 300, 'lr_scheduler': 'cosine', 'lr': 2e-05, 'min_lr': 1e-07, 'beta1': 0.9, 'beta2': 0.95, 'eps': 1e-15, 'ema': 0.993, 'max_grad_norm': 1.0, 'timestep_shift': 1.0, 'mse_weight': 1.0, 'ce_weight': 1.0, 'ce_loss_reweighting': False, 'expected_num_tokens': 20000, 'num_replicate': 1, 'num_shard': 8, 'sharding_strategy': 'HYBRID_SHARD', 'backward_prefetch': 'BACKWARD_PRE', 'cpu_offload': False, 'freeze_llm': False, 'freeze_vit': False, 'freeze_vae': True, 'freeze_und': False, 'copy_init_moe': True, 'use_flex': False, 'eval_every': 500, 'num_eval_batches': 20, 'use_ema_for_eval': True, 'eval_log_dir': None, 'eval_run_tag': '', 'viz_every': 500, 'viz_n': 8, 'viz_outdir': 'results/viz', 'eval_dataset_config_file': './data/configs/vlm_gym_jigsaw_train_mseloss_only.yaml', 'viz_dataset_config_file': './data/configs/vlm_gym_jigsaw_train_mseloss_only.yaml', 'eval_print_n': 3, 'save_ema_only': True, 'save_optimizer': False}
+2026-01-04 09:32:01,970 INFO    MainThread:13221 [wandb_run.py:_config_callback():1396] config_cb None None {'model_path': '/home/clouduser/Code/Models/BAGEL-7B-MoT', 'llm_path': 'hf/Qwen2.5-0.5B-Instruct/', 'llm_qk_norm': True, 'tie_word_embeddings': False, 'layer_module': 'Qwen2MoTDecoderLayer', 'vae_path': 'flux/vae/ae.safetensors', 'vit_path': 'hf/siglip-so400m-14-980-flash-attn2-navit/', 'max_latent_size': 64, 'latent_patch_size': 2, 'vit_patch_size': 14, 'vit_max_num_patch_per_side': 70, 'connector_act': 'gelu_pytorch_tanh', 'interpolate_pos': False, 'vit_select_layer': -2, 'vit_rope': False, 'text_cond_dropout_prob': 0.0, 'vae_cond_dropout_prob': 0.0, 'vit_cond_dropout_prob': 0.0}
+2026-01-04 09:32:01,971 INFO    MainThread:13221 [wandb_run.py:_config_callback():1396] config_cb None None {'dataset_config_file': './data/configs/vlm_gym_jigsaw_train_mseloss_only.yaml', 'train_data_dir': '/home/clouduser/Code/data/gym/jigsaw-swap_v5/train/', 'train_jsonl_path': '/home/clouduser/Code/data/gym/jigsaw-swap_v5/train/', 'eval_data_dir': '/home/clouduser/Code/data/gym/jigsaw-swap_v5/val/', 'eval_jsonl_path': '/home/clouduser/Code/data/gym/jigsaw-swap_v5/val/', 'inference_hash_file': '/home/clouduser/Code/Github/launch_new/hashes_test_set_v10.json', 'prefetch_factor': 2, 'num_workers': 1, 'max_num_tokens_per_sample': 20000, 'max_num_tokens': 20000, 'prefer_buffer_before': 16384, 'max_buffer_size': 50, 'data_seed': 42}