configs.yaml

defaults:

  seed: 0
  method: name
  task: dummy_disc
  logdir: ~/d3_logdir
  replay: uniform
  replay_size: 1e6
  replay_online: False
  replay_priority_p: 0.1
  replay_priority_thresh: 50
  eval_dir: ''
  filter: '.*'

  jax:
    platform: gpu
    jit: True
    precision: float16
    prealloc: True
    debug_nans: False
    logical_cpus: 0
    debug: False
    policy_devices: [0]
    train_devices: [0]
    sync_every: 10
    profiler: True
    # transfer_guard: True  # TODO: fix where it is violated.
    transfer_guard: False

  run:
    script: train_eval
    steps: 1e10
    expl_until: 0
    log_every: 900
    save_every: 7200
    eval_every: 2.5e4
    eval_initial: True
    eval_eps: 1
    eval_samples: 1
    train_ratio: 64.0
    train_fill: 0
    eval_fill: 0
    log_zeros: True
    log_keys_video: [image, fullsize_image]
    log_keys_sum: '(reward)'
    log_keys_mean: '(log_entropy)'
    log_keys_eval_report: '(openl_)'
    log_keys_max: '^$'
    log_keys_last: '^$'
    log_keys_trajectory: '^$'
    log_keys_success: '^$'
    log_video_streams: 1
    log_video_timeout: 60
    wandb: False
    wandb_entity: entity1
    wandb_project: project1
    from_checkpoint: ''
    actor_host: 'localhost'
    actor_port: '5551'
    actor_batch: 32
    actor_threads: 1
    env_replica: -1
    ipv6: False
    trace_malloc: False
    is_train_behavior: True
    rew_smoothing_mode: gaussian
    rew_smoothing_amt: 0.0

  envs: {amount: 4, parallel: process, length: 0, reset: True, restart: True, discretize: 0, checks: False}
  wrapper: {length: 0, reset: True, discretize: 0, checks: False, resize: 0}
  env:
    atari: {size: [64, 64], repeat: 4, sticky: True, gray: True, actions: all, lives: unused, noops: 0, pooling: 2, aggregate: max, resize: pillow}
    atari100k: {size: [64, 64], repeat: 4, sticky: False, gray: False, actions: needed, lives: unused, noops: 30, resize: pillow}
    dmlab: {size: [64, 64], repeat: 4, episodic: True}
    minecraft: {size: [64, 64], break_speed: 100.0, logs: False}
    dmc: {size: [64, 64], repeat: 2, camera: -1}
    loconav: {size: [64, 64], repeat: 2, camera: -1}
    agx: {height: 64, width: 64, repeat: 4}
    robodesk: {repeat: 8, length: 300, image_size: 64}
    hand: {height: 64, width: 64}

  # Agent
  task_behavior: Greedy
  expl_behavior: None
  batch_size: 16
  batch_length: 64
  eval_batch_size: 16
  eval_batch_length: 64
  data_loaders: 8

  # World Model
  grad_heads: [decoder, reward, cont]
  rssm_type: rssm
  rssm: {impl: softmax, deter: 4096, units: 1024, stoch: 32, classes: 32, act: silu, norm: layer, unimix: 0.01, unroll: False, action_clip: 1.0, bottleneck: -1, winit: normal, fan: avg, maskgit: {embed: 256, layers: 4, heads: 4, ffwdim: 256, steps: 4}}
  encoder: {mlp_keys: '.*', cnn_keys: '.*', act: silu, norm: layer, mlp_layers: 5, mlp_units: 1024, cnn: resnet, cnn_depth: 96, cnn_blocks: 0, resize: stride, winit: normal, fan: avg, symlog_inputs: True, minres: 4}
  decoder: {mlp_keys: '.*', cnn_keys: '.*', act: silu, norm: layer, mlp_layers: 5, mlp_units: 1024, cnn: resnet, cnn_depth: 96, cnn_blocks: 0, image_dist: mse, vector_dist: symlog_mse, inputs: [deter, stoch], resize: stride, winit: normal, fan: avg, outscale: 1.0, minres: 4, cnn_sigmoid: False}
  reward_head: {layers: 5, units: 1024, act: silu, norm: layer, dist: symlog_and_twohot, outscale: 0.0, outnorm: False, inputs: [deter, stoch], winit: normal, fan: avg, bins: 255}
  cont_head: {layers: 5, units: 1024, act: silu, norm: layer, dist: binary, outscale: 1.0, outnorm: False, inputs: [deter, stoch], winit: normal, fan: avg}
  loss_scales: {image: 1.0, vector: 1.0, reward: 1.0, cont: 1.0, dyn: 0.5, rep: 0.1, actor: 1.0, critic: 1.0, slowreg: 1.0}
  rssm_loss: {free: 1.0}
  model_opt: {opt: adam, lr: 1e-4, eps: 1e-8, clip: 1000.0, wd: 0.0, warmup: 0, lateclip: 0.0}

  # Actor Critic
  actor: {layers: 5, units: 1024, act: silu, norm: layer, minstd: 0.1, maxstd: 1.0, outscale: 1.0, outnorm: False, unimix: 0.01, inputs: [deter, stoch], winit: normal, fan: avg, symlog_inputs: False}
  critic: {layers: 5, units: 1024, act: silu, norm: layer, dist: symlog_and_twohot, outscale: 0.0, outnorm: False, inputs: [deter, stoch], winit: normal, fan: avg, bins: 255, symlog_inputs: False}
  actor_opt: {opt: adam, lr: 3e-5, eps: 1e-5, clip: 100.0, wd: 0.0, warmup: 0, lateclip: 0.0}
  critic_opt: {opt: adam, lr: 3e-5, eps: 1e-5, clip: 100.0, wd: 0.0, warmup: 0, lateclip: 0.0}
  actor_dist_disc: onehot
  actor_dist_cont: normal
  actor_grad_disc: reinforce
  actor_grad_cont: backprop
  critic_type: vfunction
  imag_horizon: 15
  imag_unroll: False
  imag_cont: mode
  horizon: 333
  return_lambda: 0.95
  critic_slowreg: logprob
  slow_critic_update: 1
  slow_critic_fraction: 0.02
  slow_critic_target: False
  retnorm: {impl: perc_ema, decay: 0.99, max: 1.0, perclo: 5.0, perchi: 95.0}
  actent: 3e-4

  # Exploration
  expl_rewards: {extr: 1.0, disag: 0.1}
  expl_opt: {opt: adam, lr: 1e-4, eps: 1e-5, clip: 100.0, wd: 0.0, warmup: 0}
  disag_head: {layers: 5, units: 1024, act: silu, norm: layer, dist: mse, outscale: 1.0, inputs: [deter, stoch, action], winit: normal, fan: avg}
  disag_target: [stoch]
  disag_models: 8

  # Director
  director_jointly: True
  train_skill_duration: 8
  env_skill_duration: 8
  goal_enc: {layers: 5, units: 1024, act: silu, norm: layer, dist: onehot, outscale: 1.0, inputs: [goal]}
  goal_dec: {layers: 5, units: 1024, act: silu, norm: layer, dist: mse, outscale: 0.1, inputs: [skill]}
  goal_opt: {opt: adam, lr: 1e-4, eps: 1e-6, clip: 100.0, wd: 1e-2, wd_pattern: 'kernel'}
  goal_kl_scale: 1.0
  goal_kl_free: 1.0
  skill_shape: [8, 8]
  manager_rews: {extr: 1.0, expl: 0.1, goal: 0.0}
  manager_actent: 3e-4
  worker_rews: {extr: 0.0, expl: 0.0, goal: 1.0}
  worker_actent: 3e-4
  worker_inputs: [deter, stoch, goal]
  worker_goals: [manager]
  worker_report_horizon: 64


  # Director + Dreamer v3
  advnorm: {impl: off, decay: 0.99, max: 1.0, perclo: 5.0, perchi: 95.0}  # perc_ema used for Director + Dreamer v2
  actent_norm: False
  actent_norm_cfg: {impl: mult, scale: 3e-3, target: 0.5, min: 1e-5, max: 1e2, vel: 0.1}


  # Dream High
  imag_horizon_high: 8
  wm_high: {layers: 5, units: 1024, act: silu, norm: layer, dist: symlog_mse, outscale: 0.1, inputs: [deter, action]}
  wm_high_opt: {opt: adam, lr: 1e-4, eps: 1e-6, clip: 100.0, wd: 1e-2, wd_pattern: 'kernel'}
  rew_high_head: {layers: 5, units: 1024, act: silu, norm: layer, dist: symlog_and_twohot, outscale: 0.0, outnorm: False, inputs: [deter], winit: normal, fan: avg, bins: 255}
  rew_high_opt: {opt: adam, lr: 1e-4, eps: 1e-6, clip: 100.0, wd: 1e-2, wd_pattern: 'kernel'}
  cont_high_head: {layers: 5, units: 1024, act: silu, norm: layer, dist: binary, outscale: 1.0, outnorm: False, inputs: [deter], winit: normal, fan: avg}
  cont_high_opt: {opt: adam, lr: 1e-4, eps: 1e-6, clip: 100.0, wd: 1e-2, wd_pattern: 'kernel'}
  manager_inputs: [deter]

minecraft:

  task: minecraft_diamond
  envs.amount: 16
  run:
    script: train_save
    eval_fill: 1e5
    train_ratio: 16
    log_keys_max: '^log_inventory.*'
  encoder: {mlp_keys: 'inventory|inventory_max|equipped|health|hunger|breath|reward', cnn_keys: 'image'}
  decoder: {mlp_keys: 'inventory|inventory_max|equipped|health|hunger|breath', cnn_keys: 'image'}

dmlab:

  task: dmlab_explore_goal_locations_small
  envs.amount: 8
  encoder: {mlp_keys: '$^', cnn_keys: 'image'}
  decoder: {mlp_keys: '$^', cnn_keys: 'image'}
  run.train_ratio: 64

atari:

  task: atari_pong
  envs.amount: 8
  run:
    steps: 5.5e7
    eval_eps: 10
    train_ratio: 64
  encoder: {mlp_keys: '$^', cnn_keys: 'image'}
  decoder: {mlp_keys: '$^', cnn_keys: 'image'}

atari100k:

  task: atari_pong
  envs: {amount: 1}
  run:
    script: train_eval
    steps: 1.5e5
    eval_every: 1e5
    eval_initial: False
    eval_eps: 100
    train_ratio: 1024
  jax.precision: float32
  rssm.deter: 512
  .*\.cnn_depth: 32
  .*\.layers: 2
  .*\.units$: 512
  # actor_eval_sample: True
  encoder: {mlp_keys: '$^', cnn_keys: 'image'}
  decoder: {mlp_keys: '$^', cnn_keys: 'image'}

crafter:

  task: crafter_reward
  envs.amount: 1
  run:
    log_keys_max: '^log_achievement_.*'
    log_keys_sum: '^log_reward$'
    log_keys_trajectory: '^log_achievement'
    log_keys_success: '^log_achievement'
  run.train_ratio: 512
  encoder: {mlp_keys: '$^', cnn_keys: 'image'}
  decoder: {mlp_keys: '$^', cnn_keys: 'image'}

dmc_vision:

  task: dmc_walker_walk
  run.train_ratio: 512
  rssm.deter: 512
  .*\.cnn_depth: 32
  .*\.layers: 2
  .*\.units: 512
  encoder: {mlp_keys: '$^', cnn_keys: 'image'}
  decoder: {mlp_keys: '$^', cnn_keys: 'image'}

dmc_proprio:

  task: dmc_walker_walk
  run.train_ratio: 512
  rssm.deter: 512
  .*\.cnn_depth: 32
  .*\.layers: 2
  .*\.units: 512
  encoder: {mlp_keys: '.*', cnn_keys: '$^'}
  decoder: {mlp_keys: '.*', cnn_keys: '$^'}

bsuite:

  task: bsuite_mnist/0
  envs: {amount: 1, parallel: none}
  run:
    script: train
    train_ratio: 1024  # 128 for cartpole
  rssm.deter: 512
  .*\.cnn_depth: 32
  .*\.layers: 2
  .*\.units: 512

pinpad:

  task: pinpad_six
  run.train_ratio: 32
  rssm.deter: 512
  .*\.cnn_depth: 32
  .*\.layers: 2
  .*\.units: 512
  encoder: {mlp_keys: '$^', cnn_keys: 'image'}
  decoder: {mlp_keys: '$^', cnn_keys: 'image'}

loconav:

  task: loconav_ant_maze_s
  env.loconav.repeat: 1  # repeat 1 used in Director
  run:
    train_ratio: 64
    # train_ratio: 512
    log_keys_max: '^log_.*'
  rssm.deter: 1024
  .*\.cnn_depth: 64
  .*\.layers: 4
  .*\.units: 512
  # encoder: {mlp_keys: '.*', cnn_keys: 'image'}
  # decoder: {mlp_keys: '.*', cnn_keys: 'image'}
  encoder: {mlp_keys: '.*(joints|sensors|actuator|effectors|appendages|bodies|height|zaxis).*', cnn_keys: 'image'}
  decoder: {mlp_keys: '.*(joints|sensors|actuator|effectors|appendages|bodies|height|zaxis).*', cnn_keys: 'image'}

kitchen:

  task: kitchen_kitchen-mlsh-v0
  run.train_ratio: 32
  rssm.deter: 512
  .*\.cnn_depth: 32
  .*\.layers: 2
  .*\.units: 512
  encoder: {mlp_keys: '$^', cnn_keys: 'image'}
  decoder: {mlp_keys: '$^', cnn_keys: 'image'}

agx:
  task: agx_agx-pixel-rocks-v1
  run:
    log_keys_max: '^i_.*'
    log_keys_last: '^i_.*'
    train_ratio: 64
    log_keys_video: [image1, image2, fullsize_image1, fullsize_image2]
  rssm.deter: 2048
  .*\.cnn_depth: 64
  .*\.units: 768
  .*\.layers: 4
  encoder: {mlp_keys: 'ob', cnn_keys: '(image1|image2)'}
  decoder: {mlp_keys: 'ob', cnn_keys: '(image1|image2)'}
  replay_priority_thresh: 500

agx-singleview:
  task: agx_agx-pixel-rocks-singleview-v3
  run:
    log_keys_max: '^i_.*'
    log_keys_last: '^i_.*'
    train_ratio: 64
    log_keys_video: [image1]
  rssm.deter: 512
  .*\.cnn_depth: 32
  .*\.layers: 2
  .*\.units: 512
  encoder: {mlp_keys: 'ob', cnn_keys: '(image1)'}
  decoder: {mlp_keys: 'ob', cnn_keys: '(image1)'}

robodesk:
  task: robodesk_yes
  run:
    steps: 3e6
    log_keys_max: '^i_.*'
    train_ratio: 64
  rssm.deter: 2048
  .*\.cnn_depth: 64
  .*\.units: 768
  .*\.layers: 4
  encoder: {mlp_keys: '$^', cnn_keys: 'image'}
  decoder: {mlp_keys: '$^', cnn_keys: 'image'}
  replay_priority_thresh: 50

hand:
  task: hand_RotateZ_dense
  wrapper.length: 300
  run:
    steps: 20e6
    log_keys_max: '^i_.*'
    train_ratio: 64
  rssm.deter: 2048
  .*\.cnn_depth: 64
  .*\.units: 768
  .*\.layers: 4
  encoder: {mlp_keys: '$^', cnn_keys: 'image'}
  decoder: {mlp_keys: '$^', cnn_keys: 'image'}
  replay_priority_thresh: 50

no_smoothing:
  run:
    rew_smoothing_mode: None
    rew_smoothing_amt: 0

exp_smoothing:
  run:
    rew_smoothing_mode: exponential
    rew_smoothing_amt: 0.33

gaussian_smoothing:
  run:
    rew_smoothing_mode: gaussian
    rew_smoothing_amt: 3

uniform_smoothing:
  run:
    rew_smoothing_mode: uniform
    rew_smoothing_amt: 9

uniform_before_smoothing:
  run:
    rew_smoothing_mode: uniform_before
    rew_smoothing_amt: 5

uniform_after_smoothing:
  run:
    rew_smoothing_mode: uniform_after
    rew_smoothing_amt: 5

supervised:
  # eval_batch_size: 4
  # eval_batch_length: 64
  run:
    script: train_supervised
    is_train_behavior: False
    log_every: 900
    save_every: 3600
    eval_every: 5e3

eval:
  envs.amount: 1
  wrapper.resize: 64
  run:
    script: eval_only
    log_every: 10
    steps: 10000
  env:
    atari.size: [64, 64]
    atari100k.size: [64, 64]
    dmlab.size: [64, 64]
    minecraft.size: [64, 64]
    dmc.size: [64, 64]
    loconav.size: [64, 64]
    agx:
      height: 1024
      width: 1024
    robodesk.image_size: 480
    hand:
      height: 1024
      width: 1024

morelogs:
  run:
    log_every: 10
    save_every: 20

xxxxxsmallrew:
  reward_head.layers: 1
  reward_head.units: 16

xxxxsmallrew:
  reward_head.layers: 1
  reward_head.units: 32

xxxsmallrew:
  reward_head.layers: 2
  reward_head.units: 128

xxsmallrew:
  reward_head.layers: 2
  reward_head.units: 256

xsmallrew:
  reward_head.layers: 2
  reward_head.units: 512

smallrew:
  reward_head.layers: 3
  reward_head.units: 640

largerew:
  reward_head.layers: 5
  reward_head.units: 1024

xlargerew:
  reward_head.layers: 6
  reward_head.units: 1280

small:
  rssm.deter: 512
  .*\.cnn_depth: 32
  .*\.units: 512
  .*\.layers: 2

medium:
  rssm.deter: 1024
  .*\.cnn_depth: 48
  .*\.units: 640
  .*\.layers: 3

large:
  rssm.deter: 2048
  .*\.cnn_depth: 64
  .*\.units: 768
  .*\.layers: 4

xlarge:
  rssm.deter: 4096
  .*\.cnn_depth: 96
  .*\.units: 1024
  .*\.layers: 5

multicpu:

  jax:
    logical_cpus: 8
    policy_devices: [0, 1]
    train_devices: [2, 3, 4, 5, 6, 7]
  run:
    actor_batch: 4
  envs:
    amount: 8
  batch_size: 12
  batch_length: 10

debug:

  jax: {jit: True, prealloc: False, debug: True, platform: cpu, profiler: False}
  envs: {restart: False, amount: 3}
  wrapper: {length: 100, checks: True}
  run:
    eval_every: 1000
    log_every: 5
    save_every: 10
    train_ratio: 32
    actor_batch: 2
  batch_size: 8
  batch_length: 12
  replay_size: 1e5
  encoder.cnn_depth: 8
  decoder.cnn_depth: 8
  rssm: {deter: 32, units: 16, stoch: 4, classes: 4}
  .*unroll: False
  .*\.layers: 2
  .*\.units: 16
  .*\.wd$: 0.0