domrachev03's picture
Trained SimDist world model: 24.8M steps, 90k updates, latent_dynamics 0.0278
1c5b1f4 verified
Raw
History Blame Contribute Delete
135 kB
Wrote diagnostics to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/diagnostics/step_77500
[2026-08-08 18:56:09,151][absl][INFO] - [process=0][thread=MainThread] CheckpointManager init: checkpointers=None, item_names=None, item_handlers=None, handler_registry=None
[2026-08-08 18:56:09,152][absl][INFO] - Initialized registry DefaultCheckpointHandlerRegistry({('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7619cc375f60>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7619cc375f60>}).
[2026-08-08 18:56:09,153][absl][INFO] - Cleaning up existing temporary directories at /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2.
[2026-08-08 18:56:09,155][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml) to end with ".orbax-checkpoint-tmp".
[2026-08-08 18:56:09,156][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/diagnostics could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/diagnostics) to end with ".orbax-checkpoint-tmp".
[2026-08-08 18:56:09,156][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/65000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/65000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 18:56:09,156][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/67500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/67500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 18:56:09,157][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/70000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/70000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 18:56:09,157][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/72500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/72500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 18:56:09,157][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/75000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/75000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 18:56:09,168][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml could not be identified as a finalized checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml) to be a directory.
[2026-08-08 18:56:09,185][absl][INFO] - Found 5 checkpoint steps in /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2
[2026-08-08 18:56:09,186][absl][INFO] - [process=0][thread=MainThread] CheckpointManager created, primary_host=0, CheckpointManagerOptions=CheckpointManagerOptions(save_interval_steps=1, max_to_keep=5, keep_time_interval=None, keep_period=None, should_keep_fn=None, best_fn=None, best_mode='max', keep_checkpoints_without_metrics=True, step_prefix=None, step_format_fixed_length=None, step_name_format=None, create=True, cleanup_tmp_directories=True, save_on_steps=frozenset(), single_host_load_and_broadcast=False, todelete_subdir=None, todelete_full_path=None, enable_hns=False, enable_background_delete=False, read_only=False, enable_async_checkpointing=False, async_options=None, multiprocessing_options=MultiprocessingOptions(primary_host=0, active_processes=None, barrier_sync_key_prefix=None), should_save_fn=None, file_options=FileOptions(path_permission_mode=None), save_root_metadata=True, temporary_path_class=None, save_decision_policy=None, preservation_policy=None, prevent_write_metrics=False, enable_should_save_is_saving_in_progress_check=True, enable_per_process_directory_creation=False), root_directory=/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2: <orbax.checkpoint.checkpoint_manager.CheckpointManager object at 0x7619cc376920>
[2026-08-08 18:56:09,186][absl][INFO] - [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning.
[2026-08-08 18:56:09,187][absl][INFO] - [process=0] Saving checkpoint at step 77500
[2026-08-08 18:56:09,187][absl][INFO] - [process=0] Started saving checkpoint to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500.
[2026-08-08 18:56:09,188][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500.orbax-checkpoint-tmp
[2026-08-08 18:56:09,191][absl][WARNING] - [process=0][thread=MainThread][operation_id=31] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
[2026-08-08 18:56:09,191][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp
[2026-08-08 18:56:09,192][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp
[2026-08-08 18:56:09,193][absl][INFO] - No entry found in handler registry for item: default and args with type: <class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>. Falling back to global handler registry.
[2026-08-08 18:56:09,194][absl][INFO] - Created BasePyTreeCheckpointHandler: use_ocdbt=True, use_zarr3=False, pytree_metadata_options=PyTreeMetadataOptions(support_rich_types=False), array_metadata_store=<orbax.checkpoint._src.metadata.array_metadata_store.Store object at 0x762908b3ba90>, enable_pinned_host_transfer=True, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
[2026-08-08 18:56:09,194][absl][INFO] - Deferred registration for item: "default". Adding handler `<orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler object at 0x76197c76f4c0>` for item "default" and save args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>` and restore args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardRestoreArgs'>` to `_handler_registry`.
[2026-08-08 18:56:09,194][absl][INFO] - No entry found in handler registry for item: opt_state and args with type: <class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>. Falling back to global handler registry.
[2026-08-08 18:56:09,194][absl][INFO] - Created BasePyTreeCheckpointHandler: use_ocdbt=True, use_zarr3=False, pytree_metadata_options=PyTreeMetadataOptions(support_rich_types=False), array_metadata_store=<orbax.checkpoint._src.metadata.array_metadata_store.Store object at 0x762908b3ba90>, enable_pinned_host_transfer=True, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
[2026-08-08 18:56:09,194][absl][INFO] - Deferred registration for item: "opt_state". Adding handler `<orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler object at 0x76197c76c970>` for item "opt_state" and save args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>` and restore args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardRestoreArgs'>` to `_handler_registry`.
[2026-08-08 18:56:09,261][absl][INFO] - Transferring arrays to host memory with options: use_replica_parallel=True, min_slice_bytes_for_replica_parallel=None, max_replicas_for_replica_parallel=None, enable_pinned_host_transfer=True
[2026-08-08 18:56:09,315][absl][INFO] - [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.type_handlers.ArrayHandler".serialize. Time taken: 0.054479s
[2026-08-08 18:56:09,384][absl][INFO] - Transferring arrays to host memory with options: use_replica_parallel=True, min_slice_bytes_for_replica_parallel=None, max_replicas_for_replica_parallel=None, enable_pinned_host_transfer=True
[2026-08-08 18:56:09,756][absl][INFO] - [process=0][thread=array_type_handler] Wrote 358 array_metadata.ArrayMetadata to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp/array_metadatas/process_0
[2026-08-08 18:56:09,890][absl][INFO] - [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.type_handlers.ArrayHandler".serialize. Time taken: 0.506300s
[2026-08-08 18:56:09,891][absl][INFO] - [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 0 Bytes/s (total gbytes: 45.8 MiB) (time elapsed: 694 milliseconds) (per-host)
[2026-08-08 18:56:09,891][absl][INFO] - [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 0.694951s (batch_requests_ready=0.044305s, total_serialization_initiated=0.649796s, others=0.000851s)
[2026-08-08 18:56:10,755][absl][INFO] - [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 0 Bytes/s (total gbytes: 91.6 MiB) (time elapsed: a second) (per-host)
[2026-08-08 18:56:10,768][absl][INFO] - [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 1.554464s (batch_requests_ready=0.101919s, total_serialization_initiated=0.576240s, others=0.876305s)
[2026-08-08 18:56:10,769][absl][INFO] - [process=0][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 1.578539s (all_items=0.000723s, per_item={'default': '0.00049162', 'opt_state': '0.00023127'}, temp_paths=1.577816)
[2026-08-08 18:56:10,769][absl][WARNING] - [process=0][thread=MainThread][operation_id=31] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
[2026-08-08 18:56:11,171][absl][INFO] - [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 1.274015s (commit=1.184116s, array_metadata_write=0.089900s)
[2026-08-08 18:56:11,209][absl][INFO] - [process=0] /jax/checkpoint/write/gbytes_per_sec: 0 Bytes/s (total gbytes: 45.8 MiB) (time elapsed: 2 seconds) (per-host)
[2026-08-08 18:56:11,214][absl][WARNING] - [process=0][thread=MainThread][operation_id=31] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
[2026-08-08 18:56:11,328][absl][INFO] - [process=0][thread=array_type_handler] Wrote 619 array_metadata.ArrayMetadata to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp/array_metadatas/process_0
[2026-08-08 18:56:11,720][absl][INFO] - [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 0.965552s (commit=0.885767s, array_metadata_write=0.079786s)
[2026-08-08 18:56:11,723][absl][INFO] - [process=0] /jax/checkpoint/write/gbytes_per_sec: 0 Bytes/s (total gbytes: 91.6 MiB) (time elapsed: 2 seconds) (per-host)
[2026-08-08 18:56:11,734][absl][INFO] - [process=0][thread=MainThread] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0.
[2026-08-08 18:56:11,757][absl][INFO] - [process=0][thread=MainThread] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.028502s. use_zarr3=False, enable_post_merge_validation=True, directory=/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp
[2026-08-08 18:56:11,759][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500.orbax-checkpoint-tmp/default
[2026-08-08 18:56:11,769][absl][INFO] - [process=0][thread=MainThread] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0.
[2026-08-08 18:56:11,796][absl][INFO] - [process=0][thread=MainThread] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.034410s. use_zarr3=False, enable_post_merge_validation=True, directory=/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp
[2026-08-08 18:56:11,797][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500.orbax-checkpoint-tmp/opt_state
[2026-08-08 18:56:11,800][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500
[2026-08-08 18:56:11,802][absl][INFO] - [process=0][thread=MainThread] Finished saving checkpoint (finalized tmp dir) to `/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500`.
[2026-08-08 18:56:11,803][absl][INFO] - Finished synchronous save in 2.62 seconds to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500
[2026-08-08 18:56:11,805][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml) to end with ".orbax-checkpoint-tmp".
[2026-08-08 18:56:11,806][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/diagnostics could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/diagnostics) to end with ".orbax-checkpoint-tmp".
[2026-08-08 18:56:11,806][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/65000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/65000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 18:56:11,807][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/67500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/67500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 18:56:11,807][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/70000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/70000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 18:56:11,808][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/72500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/72500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 18:56:11,808][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/75000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/75000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 18:56:11,808][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 18:56:11,815][absl][INFO] - Deleted step 65000.
[2026-08-08 18:56:11,815][absl][INFO] - [process=0][thread=MainThread][step=77500] CheckpointManager Save Finalize is syncing with other hosts...
[2026-08-08 18:56:11,815][absl][INFO] - [process=0][thread=MainThread][step=77500] CheckpointManager Save Finalize is done on all hosts.
[2026-08-08 18:56:11,816][absl][INFO] - [process=0][thread=MainThread][step=77500] Finished synchronous save.
[2026-08-08 18:56:11,816][absl][INFO] - {'step': 77500, 'event_type': 'save', 'directory': '/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': True, 'wait_for_prev_start_time': 1786215369.1869698, 'wait_for_prev_duration_secs': 8.463859558105469e-05, 'checkpointer_blocking_start_time': 1786215369.1872566, 'checkpointer_blocking_duration_secs': 2.6165382862091064, 'get_old_steps_start_time': 1786215371.803832, 'get_old_steps_duration_secs': 0.00016355514526367188, 'checkpoint_manager_blocking_start_time': 1786215369.1868665, 'checkpoint_manager_blocking_duration_secs': 2.629302978515625}
[2026-08-08 18:56:11,816][absl][INFO] - [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning.
[2026-08-08 18:56:11,816][absl][INFO] - Closing _NonBlockingMetadataStore(enable_write=True, _write_lock=<locked _thread.RLock object owner=129923580162048 count=1 at 0x7619a04fd3c0>, _store_impl=<orbax.checkpoint._src.metadata.checkpoint._MetadataStoreImpl object at 0x7619cc376c50>, _single_thread_executor=<concurrent.futures.thread.ThreadPoolExecutor object at 0x7619cc374ee0>, _write_futures=[])
Steps: 77500, Metrics: {'steps_per_second': 2.223805761615172, 'throughput/samples_per_s': 569.294274973484, 'lr': 0.00010590485180728137, 'grad_norm': 2.959945218408108, 'dataloader_wait_frac': 0.014640432691202972, 'gpu_util': 10.0, 'gpu_mem_used_frac': 0.7556966295580704, 'train/loss': 1.8115915060043335, 'train/latent_dynamics': 0.03509516268968582, 'train/reward': 0.09630102664232254, 'train/value': 0.11673177033662796, 'train/action': 1.563463568687439, 'test/loss': 0.3625366985797882, 'test/latent_dynamics': 0.03166506066918373, 'test/reward': 0.09030120074748993, 'test/value': 0.04829816520214081, 'test/action': 0.19227223098278046, 'test/dataloader_wait_frac': 0.09258119643627996, 'train/unweighted/latent_dynamics': 0.03509516268968582, 'train/unweighted/reward': 0.09630102664232254, 'train/unweighted/value': 0.11673177033662796, 'train/unweighted/action': 0.39086589217185974, 'test/unweighted/latent_dynamics': 0.03166506066918373, 'test/unweighted/reward': 0.09030120074748993, 'test/unweighted/value': 0.04829816520214081, 'test/unweighted/action': 0.048068057745695114, 'eval/latent_rollout_error/h1': 0.01902947691269219, 'eval/latent_rollout_error/h2': 0.02590755606070161, 'eval/latent_rollout_error/h3': 0.03208049386739731, 'eval/latent_rollout_error/h4': 0.03829987347126007, 'eval/latent_rollout_error/h5': 0.04329092847183347, 'eval/latent_rollout_error_mean': 0.03172166575677693, 'eval/latent_variance/encoder': 0.4032224379479885, 'eval/latent_variance/predicted': 0.3734528236091137, 'eval/reward_r2': 0.8657696363367613, 'eval/reward_explained_variance': 0.8659159601879088, 'eval/reward_pearson_r': 0.9307447939651985, 'eval/value_r2': 0.952680986841906, 'eval/value_explained_variance': 0.9537749228122561, 'eval/value_pearson_r': 0.9769653645368003, 'eval/bc_action_mse': 0.11979678725108049, 'eval/num_batches': 8.0, 'eval/value_auc_success': 0.7265625, 'eval/value_auc_num_episodes': 32.0, 'eval/value_auc_success_fraction': 0.5, 'eval/value_auc_holdout_episodewise': 0.0, 'epoch': 0}
self.pid = os.fork()
Wrote diagnostics to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/diagnostics/step_80000
[2026-08-08 19:16:27,815][absl][INFO] - [process=0][thread=MainThread] CheckpointManager init: checkpointers=None, item_names=None, item_handlers=None, handler_registry=None
[2026-08-08 19:16:27,816][absl][INFO] - Initialized registry DefaultCheckpointHandlerRegistry({('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x76197c36b760>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x76197c36b760>}).
[2026-08-08 19:16:27,817][absl][INFO] - Cleaning up existing temporary directories at /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2.
[2026-08-08 19:16:27,819][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:16:27,819][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/diagnostics could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/diagnostics) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:16:27,820][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/67500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/67500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:16:27,820][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/70000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/70000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:16:27,820][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/72500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/72500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:16:27,821][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/75000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/75000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:16:27,821][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:16:27,832][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml could not be identified as a finalized checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml) to be a directory.
[2026-08-08 19:16:27,847][absl][INFO] - Found 5 checkpoint steps in /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2
[2026-08-08 19:16:27,848][absl][INFO] - [process=0][thread=MainThread] CheckpointManager created, primary_host=0, CheckpointManagerOptions=CheckpointManagerOptions(save_interval_steps=1, max_to_keep=5, keep_time_interval=None, keep_period=None, should_keep_fn=None, best_fn=None, best_mode='max', keep_checkpoints_without_metrics=True, step_prefix=None, step_format_fixed_length=None, step_name_format=None, create=True, cleanup_tmp_directories=True, save_on_steps=frozenset(), single_host_load_and_broadcast=False, todelete_subdir=None, todelete_full_path=None, enable_hns=False, enable_background_delete=False, read_only=False, enable_async_checkpointing=False, async_options=None, multiprocessing_options=MultiprocessingOptions(primary_host=0, active_processes=None, barrier_sync_key_prefix=None), should_save_fn=None, file_options=FileOptions(path_permission_mode=None), save_root_metadata=True, temporary_path_class=None, save_decision_policy=None, preservation_policy=None, prevent_write_metrics=False, enable_should_save_is_saving_in_progress_check=True, enable_per_process_directory_creation=False), root_directory=/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2: <orbax.checkpoint.checkpoint_manager.CheckpointManager object at 0x7619a07050c0>
[2026-08-08 19:16:27,848][absl][INFO] - [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning.
[2026-08-08 19:16:27,849][absl][INFO] - [process=0] Saving checkpoint at step 80000
[2026-08-08 19:16:27,849][absl][INFO] - [process=0] Started saving checkpoint to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000.
[2026-08-08 19:16:27,850][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000.orbax-checkpoint-tmp
[2026-08-08 19:16:27,853][absl][WARNING] - [process=0][thread=MainThread][operation_id=32] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
[2026-08-08 19:16:27,854][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp
[2026-08-08 19:16:27,854][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp
[2026-08-08 19:16:27,855][absl][INFO] - No entry found in handler registry for item: default and args with type: <class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>. Falling back to global handler registry.
[2026-08-08 19:16:27,856][absl][INFO] - Created BasePyTreeCheckpointHandler: use_ocdbt=True, use_zarr3=False, pytree_metadata_options=PyTreeMetadataOptions(support_rich_types=False), array_metadata_store=<orbax.checkpoint._src.metadata.array_metadata_store.Store object at 0x762908b3ba90>, enable_pinned_host_transfer=True, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
[2026-08-08 19:16:27,856][absl][INFO] - Deferred registration for item: "default". Adding handler `<orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler object at 0x7619904ead40>` for item "default" and save args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>` and restore args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardRestoreArgs'>` to `_handler_registry`.
[2026-08-08 19:16:27,856][absl][INFO] - No entry found in handler registry for item: opt_state and args with type: <class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>. Falling back to global handler registry.
[2026-08-08 19:16:27,856][absl][INFO] - Created BasePyTreeCheckpointHandler: use_ocdbt=True, use_zarr3=False, pytree_metadata_options=PyTreeMetadataOptions(support_rich_types=False), array_metadata_store=<orbax.checkpoint._src.metadata.array_metadata_store.Store object at 0x762908b3ba90>, enable_pinned_host_transfer=True, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
[2026-08-08 19:16:27,856][absl][INFO] - Deferred registration for item: "opt_state". Adding handler `<orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler object at 0x76197c22b130>` for item "opt_state" and save args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>` and restore args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardRestoreArgs'>` to `_handler_registry`.
[2026-08-08 19:16:27,922][absl][INFO] - Transferring arrays to host memory with options: use_replica_parallel=True, min_slice_bytes_for_replica_parallel=None, max_replicas_for_replica_parallel=None, enable_pinned_host_transfer=True
[2026-08-08 19:16:27,978][absl][INFO] - [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.type_handlers.ArrayHandler".serialize. Time taken: 0.056216s
[2026-08-08 19:16:28,929][absl][INFO] - Transferring arrays to host memory with options: use_replica_parallel=True, min_slice_bytes_for_replica_parallel=None, max_replicas_for_replica_parallel=None, enable_pinned_host_transfer=True
[2026-08-08 19:16:29,259][absl][INFO] - [process=0][thread=array_type_handler] Wrote 358 array_metadata.ArrayMetadata to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp/array_metadatas/process_0
[2026-08-08 19:16:29,346][absl][INFO] - [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.type_handlers.ArrayHandler".serialize. Time taken: 0.423015s
[2026-08-08 19:16:29,363][absl][INFO] - [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 0 Bytes/s (total gbytes: 45.8 MiB) (time elapsed: a second) (per-host)
[2026-08-08 19:16:29,368][absl][INFO] - [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 1.510291s (batch_requests_ready=0.043471s, total_serialization_initiated=1.450385s, others=0.016435s)
[2026-08-08 19:16:29,388][absl][INFO] - [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 0 Bytes/s (total gbytes: 91.6 MiB) (time elapsed: a second) (per-host)
[2026-08-08 19:16:29,393][absl][INFO] - [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 1.517911s (batch_requests_ready=0.903284s, total_serialization_initiated=0.590330s, others=0.024296s)
[2026-08-08 19:16:29,394][absl][INFO] - [process=0][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 1.541591s (all_items=0.000716s, per_item={'default': '0.00048709', 'opt_state': '0.00022936'}, temp_paths=1.540875)
[2026-08-08 19:16:29,399][absl][WARNING] - [process=0][thread=MainThread][operation_id=32] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
[2026-08-08 19:16:29,895][absl][INFO] - [process=0][thread=array_type_handler] Wrote 619 array_metadata.ArrayMetadata to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp/array_metadatas/process_0
[2026-08-08 19:16:29,944][absl][INFO] - [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 0.562713s (commit=0.493628s, array_metadata_write=0.069085s)
[2026-08-08 19:16:29,950][absl][INFO] - [process=0] /jax/checkpoint/write/gbytes_per_sec: 0 Bytes/s (total gbytes: 45.8 MiB) (time elapsed: 2 seconds) (per-host)
[2026-08-08 19:16:29,954][absl][WARNING] - [process=0][thread=MainThread][operation_id=32] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
[2026-08-08 19:16:30,508][absl][INFO] - [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 1.091615s (commit=1.011686s, array_metadata_write=0.079930s)
[2026-08-08 19:16:30,511][absl][INFO] - [process=0] /jax/checkpoint/write/gbytes_per_sec: 0 Bytes/s (total gbytes: 91.6 MiB) (time elapsed: 2 seconds) (per-host)
[2026-08-08 19:16:30,522][absl][INFO] - [process=0][thread=MainThread] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0.
[2026-08-08 19:16:30,543][absl][INFO] - [process=0][thread=MainThread] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.026277s. use_zarr3=False, enable_post_merge_validation=True, directory=/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp
[2026-08-08 19:16:30,545][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000.orbax-checkpoint-tmp/default
[2026-08-08 19:16:30,553][absl][INFO] - [process=0][thread=MainThread] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0.
[2026-08-08 19:16:30,577][absl][INFO] - [process=0][thread=MainThread] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.030459s. use_zarr3=False, enable_post_merge_validation=True, directory=/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp
[2026-08-08 19:16:30,579][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000.orbax-checkpoint-tmp/opt_state
[2026-08-08 19:16:30,581][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000
[2026-08-08 19:16:30,584][absl][INFO] - [process=0][thread=MainThread] Finished saving checkpoint (finalized tmp dir) to `/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000`.
[2026-08-08 19:16:30,585][absl][INFO] - Finished synchronous save in 2.74 seconds to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000
[2026-08-08 19:16:30,587][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:16:30,588][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/diagnostics could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/diagnostics) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:16:30,588][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/67500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/67500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:16:30,589][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/70000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/70000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:16:30,589][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/72500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/72500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:16:30,589][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/75000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/75000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:16:30,590][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:16:30,590][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:16:30,597][absl][INFO] - Deleted step 67500.
[2026-08-08 19:16:30,597][absl][INFO] - [process=0][thread=MainThread][step=80000] CheckpointManager Save Finalize is syncing with other hosts...
[2026-08-08 19:16:30,598][absl][INFO] - [process=0][thread=MainThread][step=80000] CheckpointManager Save Finalize is done on all hosts.
[2026-08-08 19:16:30,598][absl][INFO] - [process=0][thread=MainThread][step=80000] Finished synchronous save.
[2026-08-08 19:16:30,598][absl][INFO] - {'step': 80000, 'event_type': 'save', 'directory': '/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': True, 'wait_for_prev_start_time': 1786216587.8488464, 'wait_for_prev_duration_secs': 8.320808410644531e-05, 'checkpointer_blocking_start_time': 1786216587.8491225, 'checkpointer_blocking_duration_secs': 2.736384630203247, 'get_old_steps_start_time': 1786216590.585538, 'get_old_steps_duration_secs': 0.00016069412231445312, 'checkpoint_manager_blocking_start_time': 1786216587.8487427, 'checkpoint_manager_blocking_duration_secs': 2.749790906906128}
[2026-08-08 19:16:30,598][absl][INFO] - [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning.
[2026-08-08 19:16:30,599][absl][INFO] - Closing _NonBlockingMetadataStore(enable_write=True, _write_lock=<locked _thread.RLock object owner=129923580162048 count=1 at 0x7619900eae00>, _store_impl=<orbax.checkpoint._src.metadata.checkpoint._MetadataStoreImpl object at 0x76197c368c10>, _single_thread_executor=<concurrent.futures.thread.ThreadPoolExecutor object at 0x76197c368fd0>, _write_futures=[])
Steps: 80000, Metrics: {'steps_per_second': 2.1858953369386933, 'throughput/samples_per_s': 559.5892062563055, 'lr': 0.00010380676394561306, 'grad_norm': 3.2097072541236877, 'dataloader_wait_frac': 0.014587272505527778, 'gpu_util': 2.0, 'gpu_mem_used_frac': 0.7556966295580704, 'train/loss': 1.997807264328003, 'train/latent_dynamics': 0.03491310775279999, 'train/reward': 0.09499596059322357, 'train/value': 0.11520768702030182, 'train/action': 1.752692699432373, 'test/loss': 0.3552946448326111, 'test/latent_dynamics': 0.02962455525994301, 'test/reward': 0.0906500443816185, 'test/value': 0.04775635525584221, 'test/action': 0.18726356327533722, 'test/dataloader_wait_frac': 0.0860506894259314, 'train/unweighted/latent_dynamics': 0.03491310775279999, 'train/unweighted/reward': 0.09499596059322357, 'train/unweighted/value': 0.11520768702030182, 'train/unweighted/action': 0.43817317485809326, 'test/unweighted/latent_dynamics': 0.02962455525994301, 'test/unweighted/reward': 0.0906500443816185, 'test/unweighted/value': 0.04775635525584221, 'test/unweighted/action': 0.046815890818834305, 'eval/latent_rollout_error/h1': 0.01721197832375765, 'eval/latent_rollout_error/h2': 0.02369445562362671, 'eval/latent_rollout_error/h3': 0.029701126739382744, 'eval/latent_rollout_error/h4': 0.03565819235518575, 'eval/latent_rollout_error/h5': 0.040975292678922415, 'eval/latent_rollout_error_mean': 0.029448209144175053, 'eval/latent_variance/encoder': 0.3933471292257309, 'eval/latent_variance/predicted': 0.36553169414401054, 'eval/reward_r2': 0.8562007289002663, 'eval/reward_explained_variance': 0.8562415056168066, 'eval/reward_pearson_r': 0.9254635251792452, 'eval/value_r2': 0.9272054893510393, 'eval/value_explained_variance': 0.9277184746148723, 'eval/value_pearson_r': 0.9664026867583859, 'eval/bc_action_mse': 0.12010589994723948, 'eval/num_batches': 8.0, 'eval/value_auc_success': 0.67578125, 'eval/value_auc_num_episodes': 32.0, 'eval/value_auc_success_fraction': 0.5, 'eval/value_auc_holdout_episodewise': 0.0, 'epoch': 0}
self.pid = os.fork()
Wrote diagnostics to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/diagnostics/step_82500
[2026-08-08 19:36:20,140][absl][INFO] - [process=0][thread=MainThread] CheckpointManager init: checkpointers=None, item_names=None, item_handlers=None, handler_registry=None
[2026-08-08 19:36:20,148][absl][INFO] - Initialized registry DefaultCheckpointHandlerRegistry({('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x76197c22a650>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x76197c22a650>}).
[2026-08-08 19:36:20,148][absl][INFO] - Cleaning up existing temporary directories at /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2.
[2026-08-08 19:36:20,150][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:36:20,151][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/diagnostics could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/diagnostics) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:36:20,151][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/70000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/70000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:36:20,152][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/72500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/72500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:36:20,152][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/75000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/75000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:36:20,152][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:36:20,153][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:36:20,165][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml could not be identified as a finalized checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml) to be a directory.
[2026-08-08 19:36:20,185][absl][INFO] - Found 5 checkpoint steps in /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2
[2026-08-08 19:36:20,186][absl][INFO] - [process=0][thread=MainThread] CheckpointManager created, primary_host=0, CheckpointManagerOptions=CheckpointManagerOptions(save_interval_steps=1, max_to_keep=5, keep_time_interval=None, keep_period=None, should_keep_fn=None, best_fn=None, best_mode='max', keep_checkpoints_without_metrics=True, step_prefix=None, step_format_fixed_length=None, step_name_format=None, create=True, cleanup_tmp_directories=True, save_on_steps=frozenset(), single_host_load_and_broadcast=False, todelete_subdir=None, todelete_full_path=None, enable_hns=False, enable_background_delete=False, read_only=False, enable_async_checkpointing=False, async_options=None, multiprocessing_options=MultiprocessingOptions(primary_host=0, active_processes=None, barrier_sync_key_prefix=None), should_save_fn=None, file_options=FileOptions(path_permission_mode=None), save_root_metadata=True, temporary_path_class=None, save_decision_policy=None, preservation_policy=None, prevent_write_metrics=False, enable_should_save_is_saving_in_progress_check=True, enable_per_process_directory_creation=False), root_directory=/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2: <orbax.checkpoint.checkpoint_manager.CheckpointManager object at 0x7620f0210a30>
[2026-08-08 19:36:20,187][absl][INFO] - [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning.
[2026-08-08 19:36:20,188][absl][INFO] - [process=0] Saving checkpoint at step 82500
[2026-08-08 19:36:20,188][absl][INFO] - [process=0] Started saving checkpoint to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/82500.
[2026-08-08 19:36:20,189][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/82500.orbax-checkpoint-tmp
[2026-08-08 19:36:20,192][absl][WARNING] - [process=0][thread=MainThread][operation_id=33] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
[2026-08-08 19:36:20,193][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/82500.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp
[2026-08-08 19:36:20,194][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/82500.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp
[2026-08-08 19:36:20,195][absl][INFO] - No entry found in handler registry for item: default and args with type: <class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>. Falling back to global handler registry.
[2026-08-08 19:36:20,196][absl][INFO] - Created BasePyTreeCheckpointHandler: use_ocdbt=True, use_zarr3=False, pytree_metadata_options=PyTreeMetadataOptions(support_rich_types=False), array_metadata_store=<orbax.checkpoint._src.metadata.array_metadata_store.Store object at 0x762908b3ba90>, enable_pinned_host_transfer=True, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
[2026-08-08 19:36:20,196][absl][INFO] - Deferred registration for item: "default". Adding handler `<orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler object at 0x76197c36a320>` for item "default" and save args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>` and restore args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardRestoreArgs'>` to `_handler_registry`.
[2026-08-08 19:36:20,196][absl][INFO] - No entry found in handler registry for item: opt_state and args with type: <class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>. Falling back to global handler registry.
[2026-08-08 19:36:20,196][absl][INFO] - Created BasePyTreeCheckpointHandler: use_ocdbt=True, use_zarr3=False, pytree_metadata_options=PyTreeMetadataOptions(support_rich_types=False), array_metadata_store=<orbax.checkpoint._src.metadata.array_metadata_store.Store object at 0x762908b3ba90>, enable_pinned_host_transfer=True, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
[2026-08-08 19:36:20,196][absl][INFO] - Deferred registration for item: "opt_state". Adding handler `<orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler object at 0x76197c3689a0>` for item "opt_state" and save args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>` and restore args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardRestoreArgs'>` to `_handler_registry`.
[2026-08-08 19:36:20,268][absl][INFO] - Transferring arrays to host memory with options: use_replica_parallel=True, min_slice_bytes_for_replica_parallel=None, max_replicas_for_replica_parallel=None, enable_pinned_host_transfer=True
[2026-08-08 19:36:20,325][absl][INFO] - [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.type_handlers.ArrayHandler".serialize. Time taken: 0.057328s
[2026-08-08 19:36:20,395][absl][INFO] - Transferring arrays to host memory with options: use_replica_parallel=True, min_slice_bytes_for_replica_parallel=None, max_replicas_for_replica_parallel=None, enable_pinned_host_transfer=True
[2026-08-08 19:36:20,863][absl][INFO] - [process=0][thread=array_type_handler] Wrote 358 array_metadata.ArrayMetadata to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/82500.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp/array_metadatas/process_0
[2026-08-08 19:36:20,937][absl][INFO] - [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.type_handlers.ArrayHandler".serialize. Time taken: 0.543068s
[2026-08-08 19:36:20,958][absl][INFO] - [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 0 Bytes/s (total gbytes: 45.8 MiB) (time elapsed: 759 milliseconds) (per-host)
[2026-08-08 19:36:20,964][absl][INFO] - [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 0.765293s (batch_requests_ready=0.049214s, total_serialization_initiated=0.699640s, others=0.016439s)
[2026-08-08 19:36:20,999][absl][INFO] - [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 0 Bytes/s (total gbytes: 91.6 MiB) (time elapsed: 779 milliseconds) (per-host)
[2026-08-08 19:36:21,000][absl][INFO] - [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 0.779819s (batch_requests_ready=0.104717s, total_serialization_initiated=0.639214s, others=0.035888s)
[2026-08-08 19:36:21,068][absl][INFO] - [process=0][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 0.876827s (all_items=0.001043s, per_item={'default': '0.00066471', 'opt_state': '0.00037813'}, temp_paths=0.875784)
[2026-08-08 19:36:21,068][absl][WARNING] - [process=0][thread=MainThread][operation_id=33] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
[2026-08-08 19:36:21,370][absl][INFO] - [process=0][thread=array_type_handler] Wrote 619 array_metadata.ArrayMetadata to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/82500.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp/array_metadatas/process_0
[2026-08-08 19:36:21,624][absl][INFO] - [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 0.621205s (commit=0.577784s, array_metadata_write=0.043421s)
[2026-08-08 19:36:21,629][absl][INFO] - [process=0] /jax/checkpoint/write/gbytes_per_sec: 0 Bytes/s (total gbytes: 45.8 MiB) (time elapsed: a second) (per-host)
[2026-08-08 19:36:21,631][absl][WARNING] - [process=0][thread=MainThread][operation_id=33] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
[2026-08-08 19:36:21,893][absl][INFO] - [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 0.893009s (commit=0.812648s, array_metadata_write=0.080361s)
[2026-08-08 19:36:21,897][absl][INFO] - [process=0] /jax/checkpoint/write/gbytes_per_sec: 0 Bytes/s (total gbytes: 91.6 MiB) (time elapsed: a second) (per-host)
[2026-08-08 19:36:21,908][absl][INFO] - [process=0][thread=MainThread] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0.
[2026-08-08 19:36:21,930][absl][INFO] - [process=0][thread=MainThread] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.027836s. use_zarr3=False, enable_post_merge_validation=True, directory=/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/82500.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp
[2026-08-08 19:36:21,931][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/82500.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/82500.orbax-checkpoint-tmp/default
[2026-08-08 19:36:21,941][absl][INFO] - [process=0][thread=MainThread] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0.
[2026-08-08 19:36:21,967][absl][INFO] - [process=0][thread=MainThread] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.032913s. use_zarr3=False, enable_post_merge_validation=True, directory=/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/82500.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp
[2026-08-08 19:36:21,969][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/82500.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/82500.orbax-checkpoint-tmp/opt_state
[2026-08-08 19:36:21,971][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/82500.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/82500
[2026-08-08 19:36:21,973][absl][INFO] - [process=0][thread=MainThread] Finished saving checkpoint (finalized tmp dir) to `/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/82500`.
[2026-08-08 19:36:21,975][absl][INFO] - Finished synchronous save in 1.79 seconds to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/82500
[2026-08-08 19:36:21,977][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:36:21,977][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/diagnostics could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/diagnostics) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:36:21,977][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/70000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/70000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:36:21,978][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/72500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/72500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:36:21,978][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/75000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/75000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:36:21,978][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:36:21,979][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:36:21,979][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/82500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/82500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:36:21,987][absl][INFO] - Deleted step 70000.
[2026-08-08 19:36:21,988][absl][INFO] - [process=0][thread=MainThread][step=82500] CheckpointManager Save Finalize is syncing with other hosts...
[2026-08-08 19:36:21,988][absl][INFO] - [process=0][thread=MainThread][step=82500] CheckpointManager Save Finalize is done on all hosts.
[2026-08-08 19:36:21,988][absl][INFO] - [process=0][thread=MainThread][step=82500] Finished synchronous save.
[2026-08-08 19:36:21,988][absl][INFO] - {'step': 82500, 'event_type': 'save', 'directory': '/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': True, 'wait_for_prev_start_time': 1786217780.1877053, 'wait_for_prev_duration_secs': 0.0001811981201171875, 'checkpointer_blocking_start_time': 1786217780.1881862, 'checkpointer_blocking_duration_secs': 1.7873730659484863, 'get_old_steps_start_time': 1786217781.9755907, 'get_old_steps_duration_secs': 0.0001556873321533203, 'checkpoint_manager_blocking_start_time': 1786217780.187465, 'checkpoint_manager_blocking_duration_secs': 1.801396369934082}
[2026-08-08 19:36:21,989][absl][INFO] - [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning.
[2026-08-08 19:36:21,989][absl][INFO] - Closing _NonBlockingMetadataStore(enable_write=True, _write_lock=<locked _thread.RLock object owner=129923580162048 count=1 at 0x7619a04e77c0>, _store_impl=<orbax.checkpoint._src.metadata.checkpoint._MetadataStoreImpl object at 0x76197c206b30>, _single_thread_executor=<concurrent.futures.thread.ThreadPoolExecutor object at 0x76197c206f50>, _write_futures=[])
Steps: 82500, Metrics: {'steps_per_second': 2.2369057112938946, 'throughput/samples_per_s': 572.647862091237, 'lr': 0.00010215354996034876, 'grad_norm': 2.6193276660084726, 'dataloader_wait_frac': 0.014656386422210654, 'gpu_util': 100.0, 'gpu_mem_used_frac': 0.7556966295580704, 'train/loss': 2.909078598022461, 'train/latent_dynamics': 0.035124197602272034, 'train/reward': 0.09476357698440552, 'train/value': 0.11448501795530319, 'train/action': 2.664710283279419, 'test/loss': 0.35246652364730835, 'test/latent_dynamics': 0.028428947553038597, 'test/reward': 0.09076973795890808, 'test/value': 0.050183434039354324, 'test/action': 0.18308432400226593, 'test/dataloader_wait_frac': 0.07403307832759255, 'train/unweighted/latent_dynamics': 0.035124197602272034, 'train/unweighted/reward': 0.09476357698440552, 'train/unweighted/value': 0.11448501795530319, 'train/unweighted/action': 0.6661775708198547, 'test/unweighted/latent_dynamics': 0.028428947553038597, 'test/unweighted/reward': 0.09076973795890808, 'test/unweighted/value': 0.050183434039354324, 'test/unweighted/action': 0.04577108100056648, 'eval/latent_rollout_error/h1': 0.01649195305071771, 'eval/latent_rollout_error/h2': 0.02295673405751586, 'eval/latent_rollout_error/h3': 0.029069016920402646, 'eval/latent_rollout_error/h4': 0.034896362805739045, 'eval/latent_rollout_error/h5': 0.03929050685837865, 'eval/latent_rollout_error_mean': 0.028540914738550784, 'eval/latent_variance/encoder': 0.3801327235996723, 'eval/latent_variance/predicted': 0.3579341322183609, 'eval/reward_r2': 0.8451192732026092, 'eval/reward_explained_variance': 0.845189836195351, 'eval/reward_pearson_r': 0.9197788819037257, 'eval/value_r2': 0.9305456669299486, 'eval/value_explained_variance': 0.9381005530969181, 'eval/value_pearson_r': 0.970501224975795, 'eval/bc_action_mse': 0.11969463255469778, 'eval/num_batches': 8.0, 'eval/value_auc_success': 0.72265625, 'eval/value_auc_num_episodes': 32.0, 'eval/value_auc_success_fraction': 0.5, 'eval/value_auc_holdout_episodewise': 0.0, 'epoch': 0}
Epoch 1/2 completed in 40625.73 seconds.
Starting epoch 2/2
Training: 0%| | 0/84564 [00:00<?, ?batch/s]
self.pid = os.fork()
Wrote diagnostics to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/diagnostics/step_85000
[2026-08-08 19:56:32,223][absl][INFO] - [process=0][thread=MainThread] CheckpointManager init: checkpointers=None, item_names=None, item_handlers=None, handler_registry=None
[2026-08-08 19:56:32,231][absl][INFO] - Initialized registry DefaultCheckpointHandlerRegistry({('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x761990255db0>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x761990255db0>}).
[2026-08-08 19:56:32,232][absl][INFO] - Cleaning up existing temporary directories at /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2.
[2026-08-08 19:56:32,235][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:56:32,235][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/diagnostics could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/diagnostics) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:56:32,235][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/72500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/72500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:56:32,236][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/75000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/75000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:56:32,236][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:56:32,236][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:56:32,236][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/82500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/82500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:56:32,246][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml could not be identified as a finalized checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml) to be a directory.
[2026-08-08 19:56:32,264][absl][INFO] - Found 5 checkpoint steps in /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2
[2026-08-08 19:56:32,265][absl][INFO] - [process=0][thread=MainThread] CheckpointManager created, primary_host=0, CheckpointManagerOptions=CheckpointManagerOptions(save_interval_steps=1, max_to_keep=5, keep_time_interval=None, keep_period=None, should_keep_fn=None, best_fn=None, best_mode='max', keep_checkpoints_without_metrics=True, step_prefix=None, step_format_fixed_length=None, step_name_format=None, create=True, cleanup_tmp_directories=True, save_on_steps=frozenset(), single_host_load_and_broadcast=False, todelete_subdir=None, todelete_full_path=None, enable_hns=False, enable_background_delete=False, read_only=False, enable_async_checkpointing=False, async_options=None, multiprocessing_options=MultiprocessingOptions(primary_host=0, active_processes=None, barrier_sync_key_prefix=None), should_save_fn=None, file_options=FileOptions(path_permission_mode=None), save_root_metadata=True, temporary_path_class=None, save_decision_policy=None, preservation_policy=None, prevent_write_metrics=False, enable_should_save_is_saving_in_progress_check=True, enable_per_process_directory_creation=False), root_directory=/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2: <orbax.checkpoint.checkpoint_manager.CheckpointManager object at 0x76199026d2d0>
[2026-08-08 19:56:32,266][absl][INFO] - [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning.
[2026-08-08 19:56:32,266][absl][INFO] - [process=0] Saving checkpoint at step 85000
[2026-08-08 19:56:32,266][absl][INFO] - [process=0] Started saving checkpoint to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/85000.
[2026-08-08 19:56:32,267][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/85000.orbax-checkpoint-tmp
[2026-08-08 19:56:32,270][absl][WARNING] - [process=0][thread=MainThread][operation_id=34] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
[2026-08-08 19:56:32,271][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/85000.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp
[2026-08-08 19:56:32,272][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/85000.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp
[2026-08-08 19:56:32,274][absl][INFO] - No entry found in handler registry for item: default and args with type: <class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>. Falling back to global handler registry.
[2026-08-08 19:56:32,274][absl][INFO] - Created BasePyTreeCheckpointHandler: use_ocdbt=True, use_zarr3=False, pytree_metadata_options=PyTreeMetadataOptions(support_rich_types=False), array_metadata_store=<orbax.checkpoint._src.metadata.array_metadata_store.Store object at 0x762908b3ba90>, enable_pinned_host_transfer=True, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
[2026-08-08 19:56:32,274][absl][INFO] - Deferred registration for item: "default". Adding handler `<orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler object at 0x761990257d30>` for item "default" and save args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>` and restore args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardRestoreArgs'>` to `_handler_registry`.
[2026-08-08 19:56:32,275][absl][INFO] - No entry found in handler registry for item: opt_state and args with type: <class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>. Falling back to global handler registry.
[2026-08-08 19:56:32,275][absl][INFO] - Created BasePyTreeCheckpointHandler: use_ocdbt=True, use_zarr3=False, pytree_metadata_options=PyTreeMetadataOptions(support_rich_types=False), array_metadata_store=<orbax.checkpoint._src.metadata.array_metadata_store.Store object at 0x762908b3ba90>, enable_pinned_host_transfer=True, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
[2026-08-08 19:56:32,275][absl][INFO] - Deferred registration for item: "opt_state". Adding handler `<orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler object at 0x7619cc2d4580>` for item "opt_state" and save args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>` and restore args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardRestoreArgs'>` to `_handler_registry`.
[2026-08-08 19:56:32,347][absl][INFO] - Transferring arrays to host memory with options: use_replica_parallel=True, min_slice_bytes_for_replica_parallel=None, max_replicas_for_replica_parallel=None, enable_pinned_host_transfer=True
[2026-08-08 19:56:32,410][absl][INFO] - [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.type_handlers.ArrayHandler".serialize. Time taken: 0.063131s
[2026-08-08 19:56:32,477][absl][INFO] - Transferring arrays to host memory with options: use_replica_parallel=True, min_slice_bytes_for_replica_parallel=None, max_replicas_for_replica_parallel=None, enable_pinned_host_transfer=True
[2026-08-08 19:56:32,876][absl][INFO] - [process=0][thread=array_type_handler] Wrote 358 array_metadata.ArrayMetadata to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/85000.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp/array_metadatas/process_0
[2026-08-08 19:56:33,808][absl][INFO] - [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.type_handlers.ArrayHandler".serialize. Time taken: 1.331778s
[2026-08-08 19:56:33,831][absl][INFO] - [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 0 Bytes/s (total gbytes: 45.8 MiB) (time elapsed: a second) (per-host)
[2026-08-08 19:56:33,837][absl][INFO] - [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 1.559492s (batch_requests_ready=0.049492s, total_serialization_initiated=1.481524s, others=0.028476s)
[2026-08-08 19:56:33,839][absl][INFO] - [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 0 Bytes/s (total gbytes: 91.6 MiB) (time elapsed: a second) (per-host)
[2026-08-08 19:56:33,851][absl][INFO] - [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 1.552192s (batch_requests_ready=0.116454s, total_serialization_initiated=1.421623s, others=0.014116s)
[2026-08-08 19:56:33,862][absl][INFO] - [process=0][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 1.592691s (all_items=0.001158s, per_item={'default': '0.00076222', 'opt_state': '0.00039577'}, temp_paths=1.591533)
[2026-08-08 19:56:33,965][absl][WARNING] - [process=0][thread=MainThread][operation_id=34] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
[2026-08-08 19:56:34,160][absl][INFO] - [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 0.334427s (commit=0.277605s, array_metadata_write=0.056822s)
[2026-08-08 19:56:34,165][absl][INFO] - [process=0] /jax/checkpoint/write/gbytes_per_sec: 0 Bytes/s (total gbytes: 45.8 MiB) (time elapsed: a second) (per-host)
[2026-08-08 19:56:34,167][absl][WARNING] - [process=0][thread=MainThread][operation_id=34] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
[2026-08-08 19:56:34,388][absl][INFO] - [process=0][thread=array_type_handler] Wrote 619 array_metadata.ArrayMetadata to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/85000.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp/array_metadatas/process_0
[2026-08-08 19:56:35,183][absl][INFO] - [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 1.343876s (commit=1.255080s, array_metadata_write=0.088797s)
[2026-08-08 19:56:35,186][absl][INFO] - [process=0] /jax/checkpoint/write/gbytes_per_sec: 0 Bytes/s (total gbytes: 91.6 MiB) (time elapsed: 2 seconds) (per-host)
[2026-08-08 19:56:35,192][absl][INFO] - Wrote Metadata={'item_handlers': {'default': 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler', 'opt_state': 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler'}, 'metrics': {}, 'performance_metrics': {}, 'init_timestamp_nsecs': 1786218992268358620, 'commit_timestamp_nsecs': None, 'custom_metadata': {}}, json={"item_handlers": {"default": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler", "opt_state": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1786218992268358620, "commit_timestamp_nsecs": null, "custom_metadata": {}} to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/85000.orbax-checkpoint-tmp/_CHECKPOINT_METADATA
[2026-08-08 19:56:35,198][absl][INFO] - [process=0][thread=MainThread] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0.
[2026-08-08 19:56:35,221][absl][INFO] - [process=0][thread=MainThread] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.028520s. use_zarr3=False, enable_post_merge_validation=True, directory=/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/85000.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp
[2026-08-08 19:56:35,223][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/85000.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/85000.orbax-checkpoint-tmp/default
[2026-08-08 19:56:35,232][absl][INFO] - [process=0][thread=MainThread] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0.
[2026-08-08 19:56:35,259][absl][INFO] - [process=0][thread=MainThread] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.034190s. use_zarr3=False, enable_post_merge_validation=True, directory=/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/85000.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp
[2026-08-08 19:56:35,261][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/85000.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/85000.orbax-checkpoint-tmp/opt_state
[2026-08-08 19:56:35,263][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/85000.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/85000
[2026-08-08 19:56:35,266][absl][INFO] - [process=0][thread=MainThread] Finished saving checkpoint (finalized tmp dir) to `/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/85000`.
[2026-08-08 19:56:35,268][absl][INFO] - Finished synchronous save in 3.00 seconds to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/85000
[2026-08-08 19:56:35,270][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:56:35,271][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/diagnostics could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/diagnostics) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:56:35,271][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/72500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/72500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:56:35,272][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/75000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/75000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:56:35,272][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:56:35,272][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:56:35,273][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/82500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/82500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:56:35,273][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/85000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/85000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 19:56:35,283][absl][INFO] - Deleted step 72500.
[2026-08-08 19:56:35,283][absl][INFO] - [process=0][thread=MainThread][step=85000] CheckpointManager Save Finalize is syncing with other hosts...
[2026-08-08 19:56:35,283][absl][INFO] - [process=0][thread=MainThread][step=85000] CheckpointManager Save Finalize is done on all hosts.
[2026-08-08 19:56:35,283][absl][INFO] - [process=0][thread=MainThread][step=85000] Finished synchronous save.
[2026-08-08 19:56:35,283][absl][INFO] - {'step': 85000, 'event_type': 'save', 'directory': '/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': True, 'wait_for_prev_start_time': 1786218992.2661204, 'wait_for_prev_duration_secs': 9.298324584960938e-05, 'checkpointer_blocking_start_time': 1786218992.2664042, 'checkpointer_blocking_duration_secs': 3.0020968914031982, 'get_old_steps_start_time': 1786218995.2685359, 'get_old_steps_duration_secs': 0.0001685619354248047, 'checkpoint_manager_blocking_start_time': 1786218992.2659879, 'checkpoint_manager_blocking_duration_secs': 3.01765775680542}
[2026-08-08 19:56:35,283][absl][INFO] - [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning.
[2026-08-08 19:56:35,284][absl][INFO] - Closing _NonBlockingMetadataStore(enable_write=True, _write_lock=<locked _thread.RLock object owner=129923580162048 count=1 at 0x7619a064b940>, _store_impl=<orbax.checkpoint._src.metadata.checkpoint._MetadataStoreImpl object at 0x761990255690>, _single_thread_executor=<concurrent.futures.thread.ThreadPoolExecutor object at 0x7619902559c0>, _write_futures=[])
Steps: 85000, Metrics: {'steps_per_second': 2.191051559261465, 'throughput/samples_per_s': 560.909199170935, 'lr': 0.00010096110781887546, 'grad_norm': 2.664278527867794, 'dataloader_wait_frac': 0.01972621496386623, 'gpu_util': 12.0, 'gpu_mem_used_frac': 0.7556966295580704, 'train/loss': 1.2923389673233032, 'train/latent_dynamics': 0.03426206484436989, 'train/reward': 0.09400884062051773, 'train/value': 0.11071474850177765, 'train/action': 1.0533535480499268, 'test/loss': 0.3499971330165863, 'test/latent_dynamics': 0.029449978843331337, 'test/reward': 0.09152733534574509, 'test/value': 0.04927626624703407, 'test/action': 0.17974364757537842, 'test/dataloader_wait_frac': 0.06768066026471253, 'train/unweighted/latent_dynamics': 0.03426206484436989, 'train/unweighted/reward': 0.09400884062051773, 'train/unweighted/value': 0.11071474850177765, 'train/unweighted/action': 0.2633383870124817, 'test/unweighted/latent_dynamics': 0.029449978843331337, 'test/unweighted/reward': 0.09152733534574509, 'test/unweighted/value': 0.04927626624703407, 'test/unweighted/action': 0.044935911893844604, 'eval/latent_rollout_error/h1': 0.017256082966923714, 'eval/latent_rollout_error/h2': 0.02419969066977501, 'eval/latent_rollout_error/h3': 0.029760430799797177, 'eval/latent_rollout_error/h4': 0.03685018629767001, 'eval/latent_rollout_error/h5': 0.041485327295958996, 'eval/latent_rollout_error_mean': 0.02991034360602498, 'eval/latent_variance/encoder': 0.3935634419322014, 'eval/latent_variance/predicted': 0.3752981536090374, 'eval/reward_r2': 0.8555056676006516, 'eval/reward_explained_variance': 0.8578352471332622, 'eval/reward_pearson_r': 0.9264707654358145, 'eval/value_r2': 0.9153592203993111, 'eval/value_explained_variance': 0.9153896076849279, 'eval/value_pearson_r': 0.9645543510214282, 'eval/bc_action_mse': 0.11925048275634284, 'eval/num_batches': 8.0, 'eval/value_auc_success': 0.66015625, 'eval/value_auc_num_episodes': 32.0, 'eval/value_auc_success_fraction': 0.5, 'eval/value_auc_holdout_episodewise': 0.0, 'epoch': 1}
self.pid = os.fork()
Wrote diagnostics to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/diagnostics/step_87500
[2026-08-08 20:15:58,354][absl][INFO] - [process=0][thread=MainThread] CheckpointManager init: checkpointers=None, item_names=None, item_handlers=None, handler_registry=None
[2026-08-08 20:15:58,356][absl][INFO] - Initialized registry DefaultCheckpointHandlerRegistry({('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7619a03dae00>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7619a03dae00>}).
[2026-08-08 20:15:58,357][absl][INFO] - Cleaning up existing temporary directories at /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2.
[2026-08-08 20:15:58,359][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml) to end with ".orbax-checkpoint-tmp".
[2026-08-08 20:15:58,359][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/diagnostics could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/diagnostics) to end with ".orbax-checkpoint-tmp".
[2026-08-08 20:15:58,359][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/75000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/75000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 20:15:58,360][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 20:15:58,360][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 20:15:58,360][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/82500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/82500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 20:15:58,360][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/85000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/85000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 20:15:58,371][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml could not be identified as a finalized checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml) to be a directory.
[2026-08-08 20:15:58,388][absl][INFO] - Found 5 checkpoint steps in /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2
[2026-08-08 20:15:58,389][absl][INFO] - [process=0][thread=MainThread] CheckpointManager created, primary_host=0, CheckpointManagerOptions=CheckpointManagerOptions(save_interval_steps=1, max_to_keep=5, keep_time_interval=None, keep_period=None, should_keep_fn=None, best_fn=None, best_mode='max', keep_checkpoints_without_metrics=True, step_prefix=None, step_format_fixed_length=None, step_name_format=None, create=True, cleanup_tmp_directories=True, save_on_steps=frozenset(), single_host_load_and_broadcast=False, todelete_subdir=None, todelete_full_path=None, enable_hns=False, enable_background_delete=False, read_only=False, enable_async_checkpointing=False, async_options=None, multiprocessing_options=MultiprocessingOptions(primary_host=0, active_processes=None, barrier_sync_key_prefix=None), should_save_fn=None, file_options=FileOptions(path_permission_mode=None), save_root_metadata=True, temporary_path_class=None, save_decision_policy=None, preservation_policy=None, prevent_write_metrics=False, enable_should_save_is_saving_in_progress_check=True, enable_per_process_directory_creation=False), root_directory=/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2: <orbax.checkpoint.checkpoint_manager.CheckpointManager object at 0x761a4dca1a80>
[2026-08-08 20:15:58,390][absl][INFO] - [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning.
[2026-08-08 20:15:58,390][absl][INFO] - [process=0] Saving checkpoint at step 87500
[2026-08-08 20:15:58,390][absl][INFO] - [process=0] Started saving checkpoint to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/87500.
[2026-08-08 20:15:58,391][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/87500.orbax-checkpoint-tmp
[2026-08-08 20:15:58,394][absl][WARNING] - [process=0][thread=MainThread][operation_id=35] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
[2026-08-08 20:15:58,395][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/87500.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp
[2026-08-08 20:15:58,395][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/87500.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp
[2026-08-08 20:15:58,397][absl][INFO] - No entry found in handler registry for item: default and args with type: <class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>. Falling back to global handler registry.
[2026-08-08 20:15:58,397][absl][INFO] - Created BasePyTreeCheckpointHandler: use_ocdbt=True, use_zarr3=False, pytree_metadata_options=PyTreeMetadataOptions(support_rich_types=False), array_metadata_store=<orbax.checkpoint._src.metadata.array_metadata_store.Store object at 0x762908b3ba90>, enable_pinned_host_transfer=True, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
[2026-08-08 20:15:58,397][absl][INFO] - Deferred registration for item: "default". Adding handler `<orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler object at 0x76197c22bb80>` for item "default" and save args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>` and restore args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardRestoreArgs'>` to `_handler_registry`.
[2026-08-08 20:15:58,397][absl][INFO] - No entry found in handler registry for item: opt_state and args with type: <class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>. Falling back to global handler registry.
[2026-08-08 20:15:58,397][absl][INFO] - Created BasePyTreeCheckpointHandler: use_ocdbt=True, use_zarr3=False, pytree_metadata_options=PyTreeMetadataOptions(support_rich_types=False), array_metadata_store=<orbax.checkpoint._src.metadata.array_metadata_store.Store object at 0x762908b3ba90>, enable_pinned_host_transfer=True, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
[2026-08-08 20:15:58,397][absl][INFO] - Deferred registration for item: "opt_state". Adding handler `<orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler object at 0x7619cc2d5870>` for item "opt_state" and save args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>` and restore args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardRestoreArgs'>` to `_handler_registry`.
[2026-08-08 20:15:58,465][absl][INFO] - Transferring arrays to host memory with options: use_replica_parallel=True, min_slice_bytes_for_replica_parallel=None, max_replicas_for_replica_parallel=None, enable_pinned_host_transfer=True
[2026-08-08 20:15:59,277][absl][INFO] - [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.type_handlers.ArrayHandler".serialize. Time taken: 0.812615s
[2026-08-08 20:15:59,345][absl][INFO] - Transferring arrays to host memory with options: use_replica_parallel=True, min_slice_bytes_for_replica_parallel=None, max_replicas_for_replica_parallel=None, enable_pinned_host_transfer=True
[2026-08-08 20:15:59,782][absl][INFO] - [process=0][thread=array_type_handler] Wrote 358 array_metadata.ArrayMetadata to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/87500.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp/array_metadatas/process_0
[2026-08-08 20:15:59,871][absl][INFO] - [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.type_handlers.ArrayHandler".serialize. Time taken: 0.527652s
[2026-08-08 20:15:59,883][absl][INFO] - [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 0 Bytes/s (total gbytes: 45.8 MiB) (time elapsed: a second) (per-host)
[2026-08-08 20:15:59,895][absl][INFO] - [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 1.495795s (batch_requests_ready=0.044896s, total_serialization_initiated=1.427759s, others=0.023140s)
[2026-08-08 20:15:59,913][absl][INFO] - [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 0 Bytes/s (total gbytes: 91.6 MiB) (time elapsed: a second) (per-host)
[2026-08-08 20:15:59,919][absl][INFO] - [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 1.501843s (batch_requests_ready=0.861057s, total_serialization_initiated=0.622729s, others=0.018056s)
[2026-08-08 20:15:59,920][absl][INFO] - [process=0][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 1.526602s (all_items=0.000696s, per_item={'default': '0.00046301', 'opt_state': '0.00023270'}, temp_paths=1.525906)
[2026-08-08 20:15:59,944][absl][WARNING] - [process=0][thread=MainThread][operation_id=35] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
[2026-08-08 20:16:00,429][absl][INFO] - [process=0][thread=array_type_handler] Wrote 619 array_metadata.ArrayMetadata to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/87500.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp/array_metadatas/process_0
[2026-08-08 20:16:00,911][absl][INFO] - [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 1.003942s (commit=0.956271s, array_metadata_write=0.047671s)
[2026-08-08 20:16:00,916][absl][INFO] - [process=0] /jax/checkpoint/write/gbytes_per_sec: 0 Bytes/s (total gbytes: 45.8 MiB) (time elapsed: 2 seconds) (per-host)
[2026-08-08 20:16:00,919][absl][WARNING] - [process=0][thread=MainThread][operation_id=35] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
[2026-08-08 20:16:01,377][absl][INFO] - [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 1.456581s (commit=1.387865s, array_metadata_write=0.068716s)
[2026-08-08 20:16:01,380][absl][INFO] - [process=0] /jax/checkpoint/write/gbytes_per_sec: 0 Bytes/s (total gbytes: 91.6 MiB) (time elapsed: 2 seconds) (per-host)
[2026-08-08 20:16:01,390][absl][INFO] - [process=0][thread=MainThread] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0.
[2026-08-08 20:16:01,412][absl][INFO] - [process=0][thread=MainThread] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.027213s. use_zarr3=False, enable_post_merge_validation=True, directory=/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/87500.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp
[2026-08-08 20:16:01,414][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/87500.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/87500.orbax-checkpoint-tmp/default
[2026-08-08 20:16:01,424][absl][INFO] - [process=0][thread=MainThread] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0.
[2026-08-08 20:16:01,452][absl][INFO] - [process=0][thread=MainThread] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.034665s. use_zarr3=False, enable_post_merge_validation=True, directory=/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/87500.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp
[2026-08-08 20:16:01,453][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/87500.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/87500.orbax-checkpoint-tmp/opt_state
[2026-08-08 20:16:01,456][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/87500.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/87500
[2026-08-08 20:16:01,459][absl][INFO] - [process=0][thread=MainThread] Finished saving checkpoint (finalized tmp dir) to `/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/87500`.
[2026-08-08 20:16:01,461][absl][INFO] - Finished synchronous save in 3.07 seconds to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/87500
[2026-08-08 20:16:01,463][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml) to end with ".orbax-checkpoint-tmp".
[2026-08-08 20:16:01,464][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/diagnostics could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/diagnostics) to end with ".orbax-checkpoint-tmp".
[2026-08-08 20:16:01,465][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/75000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/75000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 20:16:01,466][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 20:16:01,466][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 20:16:01,467][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/82500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/82500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 20:16:01,467][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/85000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/85000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 20:16:01,468][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/87500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/87500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 20:16:01,477][absl][INFO] - Deleted step 75000.
[2026-08-08 20:16:01,477][absl][INFO] - [process=0][thread=MainThread][step=87500] CheckpointManager Save Finalize is syncing with other hosts...
[2026-08-08 20:16:01,477][absl][INFO] - [process=0][thread=MainThread][step=87500] CheckpointManager Save Finalize is done on all hosts.
[2026-08-08 20:16:01,479][absl][INFO] - [process=0][thread=MainThread][step=87500] Finished synchronous save.
[2026-08-08 20:16:01,479][absl][INFO] - {'step': 87500, 'event_type': 'save', 'directory': '/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': True, 'wait_for_prev_start_time': 1786220158.3901062, 'wait_for_prev_duration_secs': 9.1552734375e-05, 'checkpointer_blocking_start_time': 1786220158.3903823, 'checkpointer_blocking_duration_secs': 3.071183919906616, 'get_old_steps_start_time': 1786220161.4616144, 'get_old_steps_duration_secs': 0.0002377033233642578, 'checkpoint_manager_blocking_start_time': 1786220158.3899915, 'checkpoint_manager_blocking_duration_secs': 3.0894412994384766}
[2026-08-08 20:16:01,480][absl][INFO] - [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning.
[2026-08-08 20:16:01,480][absl][INFO] - Closing _NonBlockingMetadataStore(enable_write=True, _write_lock=<locked _thread.RLock object owner=129923580162048 count=1 at 0x761a4df1a780>, _store_impl=<orbax.checkpoint._src.metadata.checkpoint._MetadataStoreImpl object at 0x761a40235420>, _single_thread_executor=<concurrent.futures.thread.ThreadPoolExecutor object at 0x761a402340d0>, _write_futures=[])
Steps: 87500, Metrics: {'steps_per_second': 2.2855638775837877, 'throughput/samples_per_s': 585.1043526614496, 'lr': 0.00010024095536209643, 'grad_norm': 3.9284848890066146, 'dataloader_wait_frac': 0.014968870037864667, 'gpu_util': 0.0, 'gpu_mem_used_frac': 0.7556966295580704, 'train/loss': 1.6484307050704956, 'train/latent_dynamics': 0.03258352354168892, 'train/reward': 0.09275206923484802, 'train/value': 0.13881859183311462, 'train/action': 1.3842756748199463, 'test/loss': 0.351699560880661, 'test/latent_dynamics': 0.02748098410665989, 'test/reward': 0.08842066675424576, 'test/value': 0.04794136807322502, 'test/action': 0.187856525182724, 'test/dataloader_wait_frac': 0.061164046657752824, 'train/unweighted/latent_dynamics': 0.03258352354168892, 'train/unweighted/reward': 0.09275206923484802, 'train/unweighted/value': 0.13881859183311462, 'train/unweighted/action': 0.3460689187049866, 'test/unweighted/latent_dynamics': 0.02748098410665989, 'test/unweighted/reward': 0.08842066675424576, 'test/unweighted/value': 0.04794136807322502, 'test/unweighted/action': 0.046964131295681, 'eval/latent_rollout_error/h1': 0.015983542893081903, 'eval/latent_rollout_error/h2': 0.02256822888739407, 'eval/latent_rollout_error/h3': 0.02805278543382883, 'eval/latent_rollout_error/h4': 0.03399595757946372, 'eval/latent_rollout_error/h5': 0.0383128160610795, 'eval/latent_rollout_error_mean': 0.027782666170969604, 'eval/latent_variance/encoder': 0.37536946311593056, 'eval/latent_variance/predicted': 0.3522566817700863, 'eval/reward_r2': 0.8542284460554775, 'eval/reward_explained_variance': 0.8542324616235485, 'eval/reward_pearson_r': 0.9244794870134933, 'eval/value_r2': 0.9435963573705607, 'eval/value_explained_variance': 0.9441889662994097, 'eval/value_pearson_r': 0.9723895464029698, 'eval/bc_action_mse': 0.11862212250780478, 'eval/num_batches': 8.0, 'eval/value_auc_success': 0.7578125, 'eval/value_auc_num_episodes': 32.0, 'eval/value_auc_success_fraction': 0.5, 'eval/value_auc_holdout_episodewise': 0.0, 'epoch': 1}
self.pid = os.fork()
Wrote diagnostics to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/diagnostics/step_90000
[2026-08-08 20:35:36,089][absl][INFO] - [process=0][thread=MainThread] CheckpointManager init: checkpointers=None, item_names=None, item_handlers=None, handler_registry=None
[2026-08-08 20:35:36,096][absl][INFO] - Initialized registry DefaultCheckpointHandlerRegistry({('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x761a40348850>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x761a40348850>}).
[2026-08-08 20:35:36,097][absl][INFO] - Cleaning up existing temporary directories at /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2.
[2026-08-08 20:35:36,099][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml) to end with ".orbax-checkpoint-tmp".
[2026-08-08 20:35:36,100][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/diagnostics could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/diagnostics) to end with ".orbax-checkpoint-tmp".
[2026-08-08 20:35:36,100][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 20:35:36,101][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 20:35:36,101][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/82500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/82500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 20:35:36,101][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/85000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/85000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 20:35:36,102][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/87500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/87500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 20:35:36,115][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml could not be identified as a finalized checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml) to be a directory.
[2026-08-08 20:35:36,133][absl][INFO] - Found 5 checkpoint steps in /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2
[2026-08-08 20:35:36,134][absl][INFO] - [process=0][thread=MainThread] CheckpointManager created, primary_host=0, CheckpointManagerOptions=CheckpointManagerOptions(save_interval_steps=1, max_to_keep=5, keep_time_interval=None, keep_period=None, should_keep_fn=None, best_fn=None, best_mode='max', keep_checkpoints_without_metrics=True, step_prefix=None, step_format_fixed_length=None, step_name_format=None, create=True, cleanup_tmp_directories=True, save_on_steps=frozenset(), single_host_load_and_broadcast=False, todelete_subdir=None, todelete_full_path=None, enable_hns=False, enable_background_delete=False, read_only=False, enable_async_checkpointing=False, async_options=None, multiprocessing_options=MultiprocessingOptions(primary_host=0, active_processes=None, barrier_sync_key_prefix=None), should_save_fn=None, file_options=FileOptions(path_permission_mode=None), save_root_metadata=True, temporary_path_class=None, save_decision_policy=None, preservation_policy=None, prevent_write_metrics=False, enable_should_save_is_saving_in_progress_check=True, enable_per_process_directory_creation=False), root_directory=/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2: <orbax.checkpoint.checkpoint_manager.CheckpointManager object at 0x761a4034b820>
[2026-08-08 20:35:36,134][absl][INFO] - [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning.
[2026-08-08 20:35:36,135][absl][INFO] - [process=0] Saving checkpoint at step 90000
[2026-08-08 20:35:36,135][absl][INFO] - [process=0] Started saving checkpoint to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/90000.
[2026-08-08 20:35:36,136][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/90000.orbax-checkpoint-tmp
[2026-08-08 20:35:36,139][absl][WARNING] - [process=0][thread=MainThread][operation_id=36] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
[2026-08-08 20:35:36,141][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/90000.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp
[2026-08-08 20:35:36,141][absl][INFO] - Creating tmp directory /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/90000.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp
[2026-08-08 20:35:36,143][absl][INFO] - No entry found in handler registry for item: default and args with type: <class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>. Falling back to global handler registry.
[2026-08-08 20:35:36,143][absl][INFO] - Created BasePyTreeCheckpointHandler: use_ocdbt=True, use_zarr3=False, pytree_metadata_options=PyTreeMetadataOptions(support_rich_types=False), array_metadata_store=<orbax.checkpoint._src.metadata.array_metadata_store.Store object at 0x762908b3ba90>, enable_pinned_host_transfer=True, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
[2026-08-08 20:35:36,143][absl][INFO] - Deferred registration for item: "default". Adding handler `<orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler object at 0x76197c3c6c50>` for item "default" and save args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>` and restore args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardRestoreArgs'>` to `_handler_registry`.
[2026-08-08 20:35:36,143][absl][INFO] - No entry found in handler registry for item: opt_state and args with type: <class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>. Falling back to global handler registry.
[2026-08-08 20:35:36,144][absl][INFO] - Created BasePyTreeCheckpointHandler: use_ocdbt=True, use_zarr3=False, pytree_metadata_options=PyTreeMetadataOptions(support_rich_types=False), array_metadata_store=<orbax.checkpoint._src.metadata.array_metadata_store.Store object at 0x762908b3ba90>, enable_pinned_host_transfer=True, save_concurrent_bytes: 96000000000 (89.4 GiB), restore_concurrent_bytes: 96000000000 (89.4 GiB)
[2026-08-08 20:35:36,144][absl][INFO] - Deferred registration for item: "opt_state". Adding handler `<orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler object at 0x76197c3c72e0>` for item "opt_state" and save args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardSaveArgs'>` and restore args `<class 'orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardRestoreArgs'>` to `_handler_registry`.
[2026-08-08 20:35:36,220][absl][INFO] - Transferring arrays to host memory with options: use_replica_parallel=True, min_slice_bytes_for_replica_parallel=None, max_replicas_for_replica_parallel=None, enable_pinned_host_transfer=True
[2026-08-08 20:35:36,279][absl][INFO] - [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.type_handlers.ArrayHandler".serialize. Time taken: 0.059730s
[2026-08-08 20:35:36,352][absl][INFO] - Transferring arrays to host memory with options: use_replica_parallel=True, min_slice_bytes_for_replica_parallel=None, max_replicas_for_replica_parallel=None, enable_pinned_host_transfer=True
[2026-08-08 20:35:36,837][absl][INFO] - [process=0][thread=array_type_handler] Wrote 358 array_metadata.ArrayMetadata to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/90000.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp/array_metadatas/process_0
[2026-08-08 20:35:36,930][absl][INFO] - [process=0][thread=MainThread] Initiated "orbax.checkpoint._src.serialization.type_handlers.ArrayHandler".serialize. Time taken: 0.579058s
[2026-08-08 20:35:36,948][absl][INFO] - [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 0 Bytes/s (total gbytes: 45.8 MiB) (time elapsed: 801 milliseconds) (per-host)
[2026-08-08 20:35:36,960][absl][INFO] - [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 0.813661s (batch_requests_ready=0.050444s, total_serialization_initiated=0.739317s, others=0.023899s)
[2026-08-08 20:35:37,012][absl][INFO] - [process=0] /jax/checkpoint/write/blocking_gbytes_per_sec: 0 Bytes/s (total gbytes: 91.6 MiB) (time elapsed: 843 milliseconds) (per-host)
[2026-08-08 20:35:37,063][absl][INFO] - [process=0][thread=MainThread] Initiated Pytree async_save. Time taken: 0.894897s (batch_requests_ready=0.116661s, total_serialization_initiated=0.690961s, others=0.087275s)
[2026-08-08 20:35:37,064][absl][INFO] - [process=0][thread=MainThread] Initiated CompositeCheckpointHandler.async_save. Time taken: 0.925230s (all_items=0.001171s, per_item={'default': '0.00072598', 'opt_state': '0.00044465'}, temp_paths=0.924060)
[2026-08-08 20:35:37,064][absl][WARNING] - [process=0][thread=MainThread][operation_id=36] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
[2026-08-08 20:35:37,595][absl][INFO] - [process=0][thread=array_type_handler] Wrote 619 array_metadata.ArrayMetadata to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/90000.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp/array_metadatas/process_0
[2026-08-08 20:35:37,655][absl][INFO] - [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 0.701964s (commit=0.652487s, array_metadata_write=0.049477s)
[2026-08-08 20:35:37,660][absl][INFO] - [process=0] /jax/checkpoint/write/gbytes_per_sec: 0 Bytes/s (total gbytes: 45.8 MiB) (time elapsed: a second) (per-host)
[2026-08-08 20:35:37,663][absl][WARNING] - [process=0][thread=MainThread][operation_id=36] _SignalingThread.join() waiting for signals ([]) blocking the main thread will slow down blocking save times. This is likely due to main thread calling result() on a CommitFuture.
[2026-08-08 20:35:38,754][absl][INFO] - [process=0][thread=write_metadata_after_commits] Commit + Array metadata written. Time taken: 1.658678s (commit=1.579103s, array_metadata_write=0.079574s)
[2026-08-08 20:35:38,757][absl][INFO] - [process=0] /jax/checkpoint/write/gbytes_per_sec: 0 Bytes/s (total gbytes: 91.6 MiB) (time elapsed: 2 seconds) (per-host)
[2026-08-08 20:35:38,771][absl][INFO] - [process=0][thread=MainThread] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0.
[2026-08-08 20:35:38,794][absl][INFO] - [process=0][thread=MainThread] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.029294s. use_zarr3=False, enable_post_merge_validation=True, directory=/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/90000.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp
[2026-08-08 20:35:38,796][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/90000.orbax-checkpoint-tmp/default.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/90000.orbax-checkpoint-tmp/default
[2026-08-08 20:35:38,806][absl][INFO] - [process=0][thread=MainThread] Skipped cross-host ArrayMetadata validation because only one process is found: process_index=0.
[2026-08-08 20:35:38,833][absl][INFO] - [process=0][thread=MainThread] Pytree save finalize (merge_ocdbt + ArrayMetadata validation) completed. Time taken: 0.034231s. use_zarr3=False, enable_post_merge_validation=True, directory=/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/90000.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp
[2026-08-08 20:35:38,835][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/90000.orbax-checkpoint-tmp/opt_state.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/90000.orbax-checkpoint-tmp/opt_state
[2026-08-08 20:35:38,837][absl][INFO] - Renaming /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/90000.orbax-checkpoint-tmp to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/90000
[2026-08-08 20:35:38,840][absl][INFO] - [process=0][thread=MainThread] Finished saving checkpoint (finalized tmp dir) to `/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/90000`.
[2026-08-08 20:35:38,842][absl][INFO] - Finished synchronous save in 2.71 seconds to /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/90000
[2026-08-08 20:35:38,844][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/model_config.yaml) to end with ".orbax-checkpoint-tmp".
[2026-08-08 20:35:38,845][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/diagnostics could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/diagnostics) to end with ".orbax-checkpoint-tmp".
[2026-08-08 20:35:38,846][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/77500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 20:35:38,846][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/80000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 20:35:38,847][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/82500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/82500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 20:35:38,848][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/85000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/85000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 20:35:38,848][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/87500 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/87500) to end with ".orbax-checkpoint-tmp".
[2026-08-08 20:35:38,849][absl][WARNING] - Path /root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/90000 could not be identified as a temporary checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected AtomicRenameTemporaryPath (/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2/90000) to end with ".orbax-checkpoint-tmp".
[2026-08-08 20:35:38,857][absl][INFO] - Deleted step 77500.
[2026-08-08 20:35:38,857][absl][INFO] - [process=0][thread=MainThread][step=90000] CheckpointManager Save Finalize is syncing with other hosts...
[2026-08-08 20:35:38,858][absl][INFO] - [process=0][thread=MainThread][step=90000] CheckpointManager Save Finalize is done on all hosts.
[2026-08-08 20:35:38,858][absl][INFO] - [process=0][thread=MainThread][step=90000] Finished synchronous save.
[2026-08-08 20:35:38,858][absl][INFO] - {'step': 90000, 'event_type': 'save', 'directory': '/root/simdist/simdist_uwlab/checkpoints/models/simdist_ur10e_24m_v2', 'reached_preemption': False, 'preemption_received_at': None, 'synchronous': True, 'wait_for_prev_start_time': 1786221336.1348066, 'wait_for_prev_duration_secs': 0.00012946128845214844, 'checkpointer_blocking_start_time': 1786221336.1352365, 'checkpointer_blocking_duration_secs': 2.70723557472229, 'get_old_steps_start_time': 1786221338.8425066, 'get_old_steps_duration_secs': 0.00018548965454101562, 'checkpoint_manager_blocking_start_time': 1786221336.134668, 'checkpoint_manager_blocking_duration_secs': 2.723936080932617}
[2026-08-08 20:35:38,858][absl][INFO] - [process=0][thread=MainThread][wait_until_finished] No Save Finalize thread to wait for. Returning.
[2026-08-08 20:35:38,859][absl][INFO] - Closing _NonBlockingMetadataStore(enable_write=True, _write_lock=<locked _thread.RLock object owner=129923580162048 count=1 at 0x761a40484380>, _store_impl=<orbax.checkpoint._src.metadata.checkpoint._MetadataStoreImpl object at 0x761a4034a1d0>, _single_thread_executor=<concurrent.futures.thread.ThreadPoolExecutor object at 0x761a4034a140>, _write_futures=[])
Steps: 90000, Metrics: {'steps_per_second': 2.264137356582644, 'throughput/samples_per_s': 579.6191632851569, 'lr': 9.999999747378752e-05, 'grad_norm': 2.907674429702759, 'dataloader_wait_frac': 0.014950468752680034, 'gpu_util': 97.0, 'gpu_mem_used_frac': 0.7556966295580704, 'train/loss': 1.6041066646575928, 'train/latent_dynamics': 0.034288667142391205, 'train/reward': 0.091888926923275, 'train/value': 0.1268676519393921, 'train/action': 1.3510606288909912, 'test/loss': 0.36793431639671326, 'test/latent_dynamics': 0.027841078117489815, 'test/reward': 0.08893612772226334, 'test/value': 0.05607955902814865, 'test/action': 0.19507770240306854, 'test/dataloader_wait_frac': 0.07222307166565173, 'train/unweighted/latent_dynamics': 0.034288667142391205, 'train/unweighted/reward': 0.091888926923275, 'train/unweighted/value': 0.1268676519393921, 'train/unweighted/action': 0.3377651572227478, 'test/unweighted/latent_dynamics': 0.027841078117489815, 'test/unweighted/reward': 0.08893612772226334, 'test/unweighted/value': 0.05607955902814865, 'test/unweighted/action': 0.048769425600767136, 'eval/latent_rollout_error/h1': 0.016440856154076755, 'eval/latent_rollout_error/h2': 0.02284725825302303, 'eval/latent_rollout_error/h3': 0.028287908528000116, 'eval/latent_rollout_error/h4': 0.03440041467547417, 'eval/latent_rollout_error/h5': 0.03848161967471242, 'eval/latent_rollout_error_mean': 0.028091611457057297, 'eval/latent_variance/encoder': 0.36949001252651215, 'eval/latent_variance/predicted': 0.3481830880045891, 'eval/reward_r2': 0.8640634588769904, 'eval/reward_explained_variance': 0.8640658966140121, 'eval/reward_pearson_r': 0.9299749404804377, 'eval/value_r2': 0.9040185561929274, 'eval/value_explained_variance': 0.904381994157503, 'eval/value_pearson_r': 0.9510478076222231, 'eval/bc_action_mse': 0.11875594078199576, 'eval/num_batches': 8.0, 'eval/value_auc_success': 0.734375, 'eval/value_auc_num_episodes': 32.0, 'eval/value_auc_success_fraction': 0.5, 'eval/value_auc_holdout_episodewise': 0.0, 'epoch': 1}
Stopping training after 90000 steps.