fabric:
  _target_: lightning.fabric.Fabric
  _convert_: all
  accelerator: gpu
  devices: ${ngpu}
  num_nodes: ${nodes}
  strategy:
    _target_: lightning.fabric.strategies.DDPStrategy
  precision: ${amp_precision}
  loggers:
  - _target_: lightning.fabric.loggers.TensorBoardLogger
    root_dir: ${save_dir}
  - _target_: lightning.pytorch.loggers.WandbLogger
    name: ${experiment_name}
    save_dir: ${save_dir}
    project: ${wandb.wandb_project}
    tags: ${wandb.wandb_tags}
    group: ${wandb.wandb_group}
    id: ${wandb.wandb_id}
    entity: ${wandb.wandb_entity}
    resume: ${wandb.wandb_resume}
ngpu: 1
nodes: 1
torch_deterministic: false
amp_precision: 32-true
eval_overrides:
  ngpu: 1
  nodes: 1
  fabric:
    loggers: null
  headless: false
  num_envs: 1
  agent:
    config:
      expert_model_path: null
      vae:
        noise_type: zeros
  motion_manager:
    motion_sampling:
      init_start_prob: 1.0
    reset_track:
      steps_min: 1000000
      steps_max: 1000001
  env:
    config:
      max_episode_length: 10000
      enable_height_termination: false
      mimic_early_termination: null
agent:
  _target_: protomotions.agents.masked_mimic.agent.MaskedMimic
  _recursive_: false
  config:
    model:
      _target_: protomotions.agents.masked_mimic.model.VaeDeterministicOutputModel
      _recursive_: false
      config:
        vae_latent_dim: ${agent.config.vae.latent_dim}
        vae_noise_type: ${agent.config.vae.noise_type}
        trunk:
          _target_: protomotions.agents.common.mlp.MultiHeadedMLP
          _recursive_: false
          num_out: ${robot.number_of_actions}
          config:
            input_models:
              self_obs:
                _target_: protomotions.agents.common.common.Flatten
                _recursive_: false
                num_in: ${env.config.humanoid_obs.obs_size}
                num_out: ${.num_in}
                config:
                  obs_key: self_obs
                  normalize_obs: true
                  norm_clamp_value: 5
              vae_latent:
                _target_: protomotions.agents.common.common.Flatten
                _recursive_: false
                num_in: ${agent.config.vae.latent_dim}
                num_out: ${.num_in}
                config:
                  obs_key: vae_latent
                  normalize_obs: false
                  norm_clamp_value: null
              terrain: ${agent.config.modules.mlp_terrain}
            trunk:
              _target_: protomotions.agents.common.mlp.MLP
              _recursive_: false
              num_out: ${robot.number_of_actions}
              config:
                layers:
                - units: 1024
                  activation: relu
                  use_layer_norm: false
                - units: 1024
                  activation: relu
                  use_layer_norm: false
                - units: 512
                  activation: relu
                  use_layer_norm: false
        optimizer:
          _target_: torch.optim.Adam
          lr: 2.0e-05
        encoder:
          _target_: protomotions.agents.masked_mimic.model.VaeModule
          _recursive_: false
          config:
            trunk:
              _target_: protomotions.agents.common.mlp.MultiHeadedMLP
              _recursive_: false
              num_out: 512
              config:
                input_models:
                  self_obs:
                    _target_: protomotions.agents.common.common.Flatten
                    _recursive_: false
                    num_in: ${env.config.humanoid_obs.obs_size}
                    num_out: ${.num_in}
                    config:
                      obs_key: self_obs
                      normalize_obs: true
                      norm_clamp_value: 5
                  mimic_target_poses:
                    _target_: protomotions.agents.common.common.Flatten
                    _recursive_: false
                    num_in: ${eval:${env.config.mimic_target_pose.num_future_steps}*${env.config.mimic_target_pose.num_obs_per_target_pose}}
                    num_out: ${.num_in}
                    config:
                      normalize_obs: true
                      norm_clamp_value: 5
                      obs_key: mimic_target_poses
                  masked_mimic_target_bodies_masks:
                    _target_: protomotions.agents.common.common.Flatten
                    config:
                      normalize_obs: false
                      obs_clamp_value: null
                      obs_key: masked_mimic_target_bodies_masks
                    num_in: ${eval:${eval:${len:${robot.trackable_bodies_subset}}+1}*2*${env.config.masked_mimic.masked_mimic_target_pose.num_future_steps}}
                    num_out: ${.num_in}
                trunk:
                  _target_: protomotions.agents.common.mlp.MLP
                  _recursive_: false
                  num_out: 512
                  config:
                    layers:
                    - units: 1024
                      activation: relu
                      use_layer_norm: false
                    - units: 1024
                      activation: relu
                      use_layer_norm: false
                    - units: 1024
                      activation: relu
                      use_layer_norm: false
                    - units: 1024
                      activation: relu
                      use_layer_norm: false
                    - units: 1024
                      activation: relu
                      use_layer_norm: false
            mu_head: ${agent.config.modules.mlp_vae_output_head}
            logvar_head: ${agent.config.modules.mlp_vae_output_head}
        prior:
          _target_: protomotions.agents.masked_mimic.model.VaeModule
          _recursive_: false
          config:
            trunk:
              _target_: protomotions.agents.common.transformer.Transformer
              _recursive_: false
              num_out: ${robot.number_of_actions}
              config:
                transformer_token_size: ${.latent_dim}
                latent_dim: ${agent.config.modules.transformer_latent_dim}
                ff_size: 1024
                num_layers: 4
                num_heads: 4
                dropout: 0
                activation: relu
                use_layer_norm: false
                input_models:
                  obs_mlp:
                    _target_: protomotions.agents.common.mlp.MLP_WithNorm
                    _recursive_: false
                    num_in: ${env.config.humanoid_obs.obs_size}
                    num_out: ${...transformer_token_size}
                    config:
                      mask_key: null
                      obs_key: self_obs
                      normalize_obs: true
                      norm_clamp_value: 5
                      layers:
                      - units: ${agent.config.modules.transformer_encoder_widths}
                        activation: relu
                        use_layer_norm: false
                      - units: ${agent.config.modules.transformer_encoder_widths}
                        activation: relu
                        use_layer_norm: false
                  masked_mimic_target_poses: ${agent.config.modules.masked_mimic_target_pose_model}
                  historical_pose_obs: ${agent.config.modules.historical_pose_obs_model}
                  motion_text_embeddings: ${agent.config.modules.motion_text_embeddings_for_transformer_model}
            mu_head: ${agent.config.modules.transformer_vae_output_head}
            logvar_head: ${agent.config.modules.transformer_vae_output_head}
    gamma: 0.99
    tau: 0.95
    e_clip: 0.2
    normalize_values: false
    normalize_advantage: false
    gradient_clip_val: 50.0
    clip_critic_loss: true
    eval_metric_keys:
    - cartesian_err
    - gt_err
    - dv_rew
    - kb_rew
    - lr_rew
    - rv_rew
    - rav_rew
    - gr_err
    - gr_err_degrees
    - masked_gt_err
    - masked_gr_err
    - masked_gr_err_degrees
    eval_length: null
    eval_num_episodes: 1
    training_early_termination:
      early_terminate_cart_err: null
      early_terminate_success_rate: null
    num_steps: 32
    fail_on_bad_grads: false
    check_grad_mag: true
    batch_size: 8192
    task_reward_w: 1.0
    num_mini_epochs: 6
    max_eval_steps: null
    eval_metrics_every: 200
    num_games: null
    manual_save_every: 10
    max_epochs: ${eval:${training_max_steps}//${ngpu}//${num_envs}//${.num_steps}}
    expert_model_path: null
    vae:
      noise_type: normal
      latent_dim: 64
      kld_schedule:
        init_kld_coeff: 0.0001
        end_kld_coeff: 0.01
        start_epoch: 3000
        end_epoch: 6000
    extra_inputs:
      mimic_target_poses: true
      masked_mimic_target_poses: true
      masked_mimic_target_bodies_masks: true
      masked_mimic_target_poses_masks: true
      motion_text_embeddings: true
      motion_text_embeddings_mask: true
      historical_pose_obs: true
      terrain: true
    modules:
      mlp_vae_output_head:
        _target_: protomotions.agents.common.mlp.MLP
        _recursive_: false
        num_in: 512
        num_out: ${agent.config.vae.latent_dim}
        config:
          layers:
          - units: 256
            activation: relu
            use_layer_norm: false
          - units: 128
            activation: relu
            use_layer_norm: false
      transformer_latent_dim: 512
      transformer_encoder_widths: 256
      transformer_vae_output_head:
        _target_: protomotions.agents.common.mlp.MLP
        _recursive_: false
        num_in: ${agent.config.model.config.prior.config.trunk.config.transformer_token_size}
        num_out: ${agent.config.vae.latent_dim}
        config:
          layers:
          - units: 256
            activation: relu
            use_layer_norm: false
          - units: 128
            activation: relu
            use_layer_norm: false
      masked_mimic_target_pose_model:
        _target_: protomotions.agents.common.mlp.MLP_WithNorm
        _recursive_: false
        num_in: ${env.config.masked_mimic.masked_mimic_target_pose.num_obs_per_sparse_target_pose}
        num_out: ${agent.config.model.config.prior.config.trunk.config.transformer_token_size}
        config:
          obs_key: masked_mimic_target_poses
          mask_key: masked_mimic_target_poses_masks
          normalize_obs: true
          norm_clamp_value: 5
          operations:
          - type: reshape
            new_shape:
            - -1
            - ${env.config.masked_mimic.masked_mimic_target_pose.num_obs_per_sparse_target_pose}
          - type: encode
          - type: reshape
            new_shape:
            - batch_size
            - ${eval:${env.config.masked_mimic.masked_mimic_target_pose.num_future_steps}+1}
            - ${agent.config.model.config.prior.config.trunk.config.transformer_token_size}
          layers:
          - units: ${agent.config.modules.transformer_encoder_widths}
            activation: relu
            use_layer_norm: false
          - units: ${agent.config.modules.transformer_encoder_widths}
            activation: relu
            use_layer_norm: false
      historical_pose_obs_model:
        _target_: protomotions.agents.common.mlp.MLP_WithNorm
        _recursive_: false
        num_in: ${eval:${env.config.humanoid_obs.obs_size}+1}
        num_out: ${agent.config.model.config.prior.config.trunk.config.transformer_token_size}
        config:
          obs_key: historical_pose_obs
          mask_key: null
          normalize_obs: true
          norm_clamp_value: 5
          operations:
          - type: reshape
            new_shape:
            - -1
            - ${eval:${env.config.humanoid_obs.obs_size}+1}
          - type: encode
          - type: reshape
            new_shape:
            - batch_size
            - ${env.config.masked_mimic.historical_obs.num_historical_conditioned_steps}
            - ${agent.config.model.config.prior.config.trunk.config.transformer_token_size}
          layers:
          - units: ${agent.config.modules.transformer_encoder_widths}
            activation: relu
            use_layer_norm: false
          - units: ${agent.config.modules.transformer_encoder_widths}
            activation: relu
            use_layer_norm: false
      motion_text_embeddings_for_transformer_model:
        _target_: protomotions.agents.common.mlp.MLP_WithNorm
        _recursive_: false
        num_in: ${env.config.masked_mimic.motion_text_embeddings.embedding_dim}
        num_out: ${agent.config.model.config.prior.config.trunk.config.transformer_token_size}
        config:
          obs_key: motion_text_embeddings
          mask_key: motion_text_embeddings_mask
          normalize_obs: true
          norm_clamp_value: 5
          operations:
          - type: encode
          - type: reshape
            new_shape:
            - batch_size
            - 1
            - ${agent.config.modules.motion_text_embeddings_for_transformer_model.num_out}
          layers:
          - units: ${agent.config.modules.transformer_encoder_widths}
            activation: relu
            use_layer_norm: false
          - units: ${agent.config.modules.transformer_encoder_widths}
            activation: relu
            use_layer_norm: false
      mlp_terrain:
        _target_: protomotions.agents.common.common.Flatten
        num_in: ${terrain.config.terrain_obs_num_samples}
        num_out: ${.num_in}
        config:
          obs_key: terrain
          normalize_obs: true
          norm_clamp_value: 5
env:
  _target_: protomotions.envs.mimic.env.Mimic
  _recursive_: false
  config:
    experiment_name: ${experiment_name}
    num_envs: ${num_envs}
    headless: ${headless}
    simulator: ${simulator}
    max_episode_length: 1000
    robot: ${robot}
    humanoid_obs:
      use_max_coords_obs: true
      obs_size: ${robot.self_obs_size}
      local_root_obs: true
      root_height_obs: true
      num_historical_steps: 120
    termination_height: 0.15
    head_termination_height: 0.3
    enable_height_termination: false
    motion_lib: ${motion_lib}
    motion_manager: ${motion_manager}
    hybrid_init_prob: 0.0
    state_init: Data
    ref_respawn_offset: ${ref_respawn_offset}
    terrain: ${terrain}
    sync_motion: ${sync_motion}
    scenes: null
    mimic_residual_control: false
    mimic_early_termination:
    - mimic_early_termination_key: max_joint_err
      mimic_early_termination_thresh: 0.5
      mimic_early_termination_thresh_on_flat: 0.25
      less_than: false
    mimic_reward_config:
      positive_constant: 0
      relative_kb_pos: false
      add_rr_to_lr: true
      rt_ignore_height: true
      mean_before_exp: true
      component_weights:
        gr_rew_w: 0.3
        lr_rew_w: 0
        gt_rew_w: 0.5
        rt_rew_w: 0
        rv_rew_w: 0
        rav_rew_w: 0
        gv_rew_w: 0.1
        gav_rew_w: 0.1
        kb_rew_w: 0
        dv_rew_w: 0.0
        rh_rew_w: 0.2
        pow_rew_w: 1.0e-05
      component_coefficients:
        gt_rew_c: -100
        rt_rew_c: -120
        rv_rew_c: -1
        rav_rew_c: -0.3
        gv_rew_c: -0.5
        gav_rew_c: -0.1
        kb_rew_c: -10
        gr_rew_c: -10
        lr_rew_c: -2
        dv_rew_c: -0.005
        rh_rew_c: -100
        kbf_rew_c: -0.01
    mimic_phase_obs:
      enabled: false
    mimic_target_pose:
      enabled: true
      type: max-coords-future-rel
      with_time: true
      with_contacts: false
      num_future_steps: 15
      num_obs_per_target_pose: ${.base_num_obs_per_target_pose}
      base_num_obs_per_target_pose: ${eval:${.num_obs_per_joint}*${robot.num_bodies}+1*${.with_time}}
      num_obs_per_joint: 18
    masked_mimic:
      enabled: true
      masked_mimic_masking:
        joint_masking:
          masked_mimic_time_gap_probability: 0.1
          time_gap_mask_min_steps: 1
          time_gap_mask_max_steps: ${eval:${env.config.mimic_target_pose.num_future_steps}+1}
          with_conditioning_time_gap_mask_max_steps: ${eval:${.time_gap_mask_max_steps}*4}
          with_conditioning_max_gap_probability: 0.2
          masked_mimic_repeat_mask_probability: 0.98
          masked_mimic_fixed_conditioning: null
          force_max_conditioned_bodies_prob: 0.1
          force_small_num_conditioned_bodies_prob: 0.1
        target_pose_visible_prob: 0.2
        motion_text_embeddings_visible_prob: 0.5
      masked_mimic_target_pose:
        num_obs_per_sparse_target_pose: ${eval:${.num_obs_per_target_pose}*${eval:${len:${robot.trackable_bodies_subset}}+1}//${robot.num_bodies}+${eval:${len:${robot.trackable_bodies_subset}}+1}*2+2}
        num_future_steps: 10
        num_obs_per_joint: 24
        num_obs_per_target_pose: ${eval:${.num_obs_per_joint}*${robot.num_bodies}}
      motion_text_embeddings:
        embedding_dim: 512
      historical_obs:
        num_historical_conditioned_steps: 15
robot:
  dof_obs_size: 138
  number_of_actions: 69
  self_obs_size: ${.self_obs_max_coords_size}
  self_obs_max_coords_size: 358
  num_bodies: ${len:${.body_names}}
  contact_bodies: ${.body_names}
  body_names:
  - Pelvis
  - L_Hip
  - L_Knee
  - L_Ankle
  - L_Toe
  - R_Hip
  - R_Knee
  - R_Ankle
  - R_Toe
  - Torso
  - Spine
  - Chest
  - Neck
  - Head
  - L_Thorax
  - L_Shoulder
  - L_Elbow
  - L_Wrist
  - L_Hand
  - R_Thorax
  - R_Shoulder
  - R_Elbow
  - R_Wrist
  - R_Hand
  dof_names:
  - L_Hip_x
  - L_Hip_y
  - L_Hip_z
  - L_Knee_x
  - L_Knee_y
  - L_Knee_z
  - L_Ankle_x
  - L_Ankle_y
  - L_Ankle_z
  - L_Toe_x
  - L_Toe_y
  - L_Toe_z
  - R_Hip_x
  - R_Hip_y
  - R_Hip_z
  - R_Knee_x
  - R_Knee_y
  - R_Knee_z
  - R_Ankle_x
  - R_Ankle_y
  - R_Ankle_z
  - R_Toe_x
  - R_Toe_y
  - R_Toe_z
  - Torso_x
  - Torso_y
  - Torso_z
  - Spine_x
  - Spine_y
  - Spine_z
  - Chest_x
  - Chest_y
  - Chest_z
  - Neck_x
  - Neck_y
  - Neck_z
  - Head_x
  - Head_y
  - Head_z
  - L_Thorax_x
  - L_Thorax_y
  - L_Thorax_z
  - L_Shoulder_x
  - L_Shoulder_y
  - L_Shoulder_z
  - L_Elbow_x
  - L_Elbow_y
  - L_Elbow_z
  - L_Wrist_x
  - L_Wrist_y
  - L_Wrist_z
  - L_Hand_x
  - L_Hand_y
  - L_Hand_z
  - R_Thorax_x
  - R_Thorax_y
  - R_Thorax_z
  - R_Shoulder_x
  - R_Shoulder_y
  - R_Shoulder_z
  - R_Elbow_x
  - R_Elbow_y
  - R_Elbow_z
  - R_Wrist_x
  - R_Wrist_y
  - R_Wrist_z
  - R_Hand_x
  - R_Hand_y
  - R_Hand_z
  trackable_bodies_subset:
  - Pelvis
  - L_Ankle
  - R_Ankle
  - L_Hand
  - R_Hand
  - Head
  dof_body_ids:
  - 1
  - 2
  - 3
  - 4
  - 5
  - 6
  - 7
  - 8
  - 9
  - 10
  - 11
  - 12
  - 13
  - 14
  - 15
  - 16
  - 17
  - 18
  - 19
  - 20
  - 21
  - 22
  - 23
  dof_effort_limits: null
  dof_vel_limits: null
  dof_armatures: null
  dof_joint_frictions: null
  key_bodies:
  - R_Ankle
  - L_Ankle
  - L_Hand
  - R_Hand
  non_termination_contact_bodies:
  - R_Ankle
  - L_Ankle
  - R_Toe
  - L_Toe
  left_foot_name: L_Ankle
  right_foot_name: R_Ankle
  head_body_name: Head
  init_state: null
  contact_pairs_multiplier: 16
  num_key_bodies: ${len:${robot.key_bodies}}
  mimic_small_marker_bodies: null
  control:
    control_type: built_in_pd
    use_biased_controller: false
    map_actions_to_pd_range: true
    stiffness: null
    damping: null
    action_scale: 1.0
    clamp_actions: 1.0
  asset:
    collapse_fixed_joints: null
    replace_cylinder_with_capsule: null
    flip_visual_attachments: null
    armature: null
    thickness: null
    max_angular_velocity: null
    max_linear_velocity: null
    density: null
    angular_damping: null
    linear_damping: null
    disable_gravity: null
    fix_base_link: null
    default_dof_drive_mode: 1
    robot_type: smpl_humanoid
    asset_root: protomotions/data/assets
    self_collisions: true
    filter_ints:
    - 0
    - 0
    - 7
    - 16
    - 12
    - 0
    - 56
    - 2
    - 33
    - 128
    - 0
    - 192
    - 0
    - 64
    - 0
    - 0
    - 0
    - 0
    - 0
    - 0
    - 0
    - 0
    - 0
    - 0
    asset_file_name: mjcf/smpl_humanoid.xml
    usd_asset_file_name: usd/smpl_humanoid.usda
  sim:
    isaacgym:
      fps: 60
      decimation: 2
      substeps: 2
    isaaclab:
      fps: 120
      decimation: 4
    genesis:
      fps: 60
      decimation: 2
      substeps: 2
simulator:
  _target_: protomotions.simulator.isaaclab.simulator.IsaacLabSimulator
  _config_target_: protomotions.simulator.isaaclab.config.IsaacLabSimulatorConfig
  _recursive_: false
  config:
    w_last: false
    headless: ${headless}
    robot: ${robot}
    num_envs: ${num_envs}
    plane:
      static_friction: 1.0
      dynamic_friction: 1.0
      restitution: 0.0
    sim:
      fps: ${robot.sim.isaaclab.fps}
      decimation: ${robot.sim.isaaclab.decimation}
      physx:
        num_threads: 4
        solver_type: 1
        num_position_iterations: 4
        num_velocity_iterations: 0
        contact_offset: 0.02
        rest_offset: 0.0
        bounce_threshold_velocity: 0.2
        max_depenetration_velocity: 10.0
        default_buffer_size_multiplier: 10.0
        gpu_found_lost_pairs_capacity: ${eval:2**21}
        gpu_max_rigid_contact_count: ${eval:2**23}
        gpu_found_lost_aggregate_pairs_capacity: ${eval:2**25}
    camera: null
    record_viewer: false
    viewer_record_dir: output/recordings/viewer
    experiment_name: ${experiment_name}
motion_lib:
  _target_: protomotions.utils.motion_lib.MotionLib
  motion_file: ${motion_file}
  ref_height_adjust: 0.0
  fix_motion_heights: true
motion_file: ???
terrain:
  _target_: protomotions.envs.base_env.env_utils.terrains.terrain.Terrain
  config:
    terrain_composition: curriculum
    map_length: 20.0
    map_width: 20.0
    border_size: 40.0
    num_levels: 7
    num_terrains: 7
    terrain_proportions:
    - 0.2
    - 0.1
    - 0.1
    - 0.1
    - 0.05
    - 0.0
    - 0.0
    - 0.45
    slope_threshold: 0.9
    num_samples_per_axis: 16
    sample_width: 1
    terrain_obs_num_samples: ${eval:${.num_samples_per_axis}**2}
    horizontal_scale: 0.1
    vertical_scale: 0.005
    spacing_between_scenes: 10
    minimal_humanoid_spacing: 0
    terrain_path: null
    load_terrain: false
    save_terrain: false
seed: 0
training_max_steps: 100000000000
codebase_version: 1.0
headless: true
sync_motion: false
checkpoint: null
experiment_name: maskedmimic
base_dir: results
save_dir: ${base_dir}/${experiment_name}
num_envs: 2048
motion_manager:
  fixed_motion_per_env: false
  fixed_motion_id: null
  motion_index_offset: null
  motion_sampling:
    init_start_prob: 0.2
  reset_track:
    steps_min: 10
    steps_max: 600
    grace_period: 5
ref_respawn_offset: 0.05
wandb:
  wandb_project: physical_animation
  wandb_tags: null
  wandb_group: null
  wandb_id: null
  wandb_entity: null
  wandb_resume: allow