# Copyright (c) 2021, salesforce.com, inc.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# For full license text, see the LICENSE file in the repo root
# or https://opensource.org/licenses/BSD-3-Clause

# YAML configuration for the tag continuous environment
name: "covid_and_economy_environment"
# Environment settings
env:
    collate_agent_step_and_reset_data: True
    components:
        - ControlUSStateOpenCloseStatus: 
            action_cooldown_period: 28
        - FederalGovernmentSubsidy:
            num_subsidy_levels: 20
            subsidy_interval: 90
            max_annual_subsidy_per_person: 20000
        - VaccinationCampaign:
            daily_vaccines_per_million_people: 3000
            delivery_interval: 1
            vaccine_delivery_start_date: "2021-01-12"
    economic_reward_crra_eta: 2
    episode_length: 540
    flatten_masks: True
    flatten_observations: False
    health_priority_scaling_agents: 0.3
    health_priority_scaling_planner: 0.45
    infection_too_sick_to_work_rate: 0.1
    multi_action_mode_agents: False
    multi_action_mode_planner: False
    n_agents: 51
    path_to_data_and_fitted_params: ""
    pop_between_age_18_65: 0.6
    risk_free_interest_rate: 0.03
    world_size: [1, 1]
    start_date: "2020-03-22"
    use_real_world_data: False
    use_real_world_policies: False
# Trainer settings
trainer:
    num_envs: 60 # number of environment replicas
    num_episodes: 1000 # number of episodes to run the training for
    train_batch_size: 5400 # total batch size used for training per iteration (across all the environments)
# Policy network settings
policy: # list all the policies below
    a:
        to_train: True # flag indicating whether the model needs to be trained
        algorithm: "PPO" # algorithm used to train the policy
        vf_loss_coeff: 1 # loss coefficient schedule for the value function loss
        entropy_coeff: 0.05 # loss coefficient schedule for the entropy loss
        gamma: 0.98 # discount factor
        lr: 0.0001 # learning rate
        model:
            type: "fully_connected"
            fc_dims: [256, 256]
            model_ckpt_filepath: ""
    p:
        to_train: True
        algorithm: "PPO"
        vf_loss_coeff: 1
        entropy_coeff: # annealing entropy over time
        - [0, 0.5]
        - [50000000, 0.05]
        gamma: 0.98
        lr: 0.0001
        model:
            type: "fully_connected"
            fc_dims: [256, 256]
            model_ckpt_filepath: ""
# Checkpoint saving setting
saving:
    metrics_log_freq: 100 # How often (in iterations) to print the metrics
    model_params_save_freq: 500 # How often (in iterations) to save the model parameters
    basedir: "/tmp" # base folder used for saving
    name: "covid19_and_economy" # experiment name
    tag: "experiments" # experiment tag
