CraftaxTrainLoop(
                study_name='craftax',
                experiment_name='exp000',
                base_dir='/opt/scratch',
                working_dir=PosixPath('/opt/scratch/runs/craftax/exp000'),
                doc='',
                step=CraftaxTrainStep.Config(
                │       model=ActorCritic.Config(
                │               observation_size=8_268,
                │               num_actions=43,
                │               channels_in=512,
                │               num_layers=3
                │       ),
                │       optimizer=PartialConfig(<class 'functools.partial'>, <class 'torch.optim.adamw.AdamW'>, lr=0.001, betas=(0.9, 0.999), weight_decay=0.01),
                │       loss=SimpleLoss.Config(
                │               loss_fn=<function priml.loss.simple_loss.bce_with_logits at 0xdefacedeface>,
                │               target_key='label',
                │               kwargs={'reduction': 'none'},
                │               channels_out=-1
                │       ),
                │       accumulate_grad_batches=1,
                │       drop_partial_accumulation_on_epoch_end=True,
                │       train_budget_steps=inf,
                │       train_budget_sec=inf,
                │       train_budget_epochs=inf,
                │       lr_schedule=PartialConfig(<class 'functools.partial'>, <function priml.math.schedules.constant at 0xdefacedeface>),
                │       parallelism=NoParallel.Config(device='auto'),
                │       model_quantization=NoModelQuantization.Config(),
                │       activation_memoization=DefaultActivationStorage.Config(),
                │       compile=PartialConfig(<class 'functools.partial'>, <function torch.compile at 0xdefacedeface>, fullgraph=True),
                │       ema=NoEMA.Config(),
                │       gradient_clip_norm=inf,
                │       skip_step_on_nonfinite_grad=False,
                │       device_init='eager',
                │       dtype_autocast=None,
                │       autocast_cache_enabled=False,
                │       env=CraftaxEnv.Config(
                │               num_envs=256,
                │               device='auto',
                │               seed=42,
                │               view=(9, 11),
                │               optimistic_reset_ratio=16,
                │               restart=RestartFromReserve.Config(),
                │               cuda_graphs=True
                │       ),
                │       rollout_steps=16,
                │       num_epochs=4,
                │       num_minibatches=8,
                │       learning_rate=0.0003,
                │       anneal_learning_rate=True,
                │       total_train_steps=244,
                │       discount=0.99,
                │       trace_decay=0.8,
                │       clip_epsilon=0.2,
                │       entropy_coefficient=0.01,
                │       value_coefficient=0.5,
                │       max_grad_norm=1.0,
                │       seed=42,
                │       cuda_graphs=True
                ),
                dataset=CraftaxRollouts.Config(updates_per_epoch=244, eval_batches=1),
                metrics_train={},
                metrics_eval={       'craftax': CraftaxScore.Config(
                                num_envs=64,
                                steps=10_000,
                                seed=42,
                                view=(9, 11),
                                device='auto'
                        )},
                checkpointer=Checkpointer.Config(
                        base_dir=PosixPath('/opt/scratch/runs/craftax/exp000'),
                        working_dir=PosixPath('/opt/scratch/runs/craftax/exp000/checkpoints'),
                        filename='step_{step:08d}.pt',
                        save_every=1_000,
                        keep_last_n=-1,
                        keep_every=0,
                        storer=SyncLocalStateDictStorer.Config(),
                        resume=True,
                        resume_step=-1,
                        allow_checkpoint_overwrite=False,
                        best_metric='',
                        best_mode='max'
                ),
                profiler=None,
                phase_timer=PhaseTimer.Config(
                        enabled=False,
                        torch_profile=False,
                        base_dir=PosixPath('/opt/scratch/runs/craftax/exp000'),
                        working_dir=PosixPath('/opt/scratch/runs/craftax/exp000/profiling'),
                        heartbeat_interval_sec=20.0,
                        fault_dump_interval_sec=40.0,
                        cuda_events=False
                ),
                tracker=None,
                max_steps=244,
                max_epochs=inf,
                max_time=inf,
                max_time_kind='wall',
                max_eval_time=inf,
                eval_stop_on_time_limit=False,
                num_steps_eval=244,
                num_steps_log=10,
                early_train_log_steps=100,
                eval_every_epoch=True,
                eval_extras_every_eval=False,
                eval_warmup_batches=0,
                eval_only=False,
                restore_rng_state=True,
                seed=None,
                num_steps_garbage_collect=inf,
                mesh_dim_model_seed='pp',
                mesh_dim_data_seed='dp',
                runtime=SingleProcess.Config(
                        device='auto',
                        deterministic=False,
                        float32_matmul_precision=None
                )
        )
