train_guidance.yaml

data_entity: two_moons
results_dir: "results/guidance/${data_entity}"
check_val_every_n_epochs: 5
num_worker: 12

max_epochs: 100
batch_size: 256
precision: 32 # 64: extended, 32: standard, 16: half

dataset:
  train_file: "data/${data_entity}_10000.pt"
  val_file: "data/${data_entity}_1000.pt"
  n_target: 20
  noise_std: 0.1 # how strong to noise the augmentations in x_target

model:
  TimeEncoder:
    enabled: True
    input_dim: 512
    output_dim: 512
    activation_func: ELU
    architecture: [512, 512,]
    final_activation: ELU # Sigmoid only needed for cosine similarity -> otherwise null
  ThetaEncoder:
    output_dim: 128
    architecture: [256]
    activation_func: "ReLU"
    final_activation: ReLU # Sigmoid only needed for cosine similarity -> otherwise null
  SimulatorEncoder: 
    output_dim: 128
    architecture: [256]
    activation_func: "ReLU"
    final_activation: ReLU # Sigmoid only needed for cosine similarity -> otherwise null
  LatentMLP:
    architecture: [256, 156, 128]
    activation_func: "ReLU"
    final_activation: Softplus   # Sigmoid only needed for cosine similarity -> otherwise null

diffusion:          
  steps: 1000
  time_repr_dim: ${model.TimeEncoder.input_dim}  # sinusoidal position encoding dimensions 
  period_spread: 10000  # parameter for sinusoidal positional encoding
  diffusion_schedule: "DDPMSchedule"
  # Variance Preserving Schedule
  VPSchedule:
    beta_start: 0.0001     
    beta_end: 0.02
    T: ${diffusion.steps}
    beta_schedule_cls: "LinearSchedule"  
  DDPMSchedule:
    beta_start: 0.0001           
    beta_end: 0.02
    T: ${diffusion.steps}
    beta_schedule_cls: "LinearSchedule"  

optimizer:
  name: AdamW
  lr: 0.0001
  weight_decay: 0.001