utils.py

import argparse
import os
import pickle
import random
import string
from datetime import datetime
from typing import List, Dict, OrderedDict, Callable, Tuple, Any, Optional

import numpy as np
import torch

from cobs import Model
from icnn import ICNN
from ppo import PPO


class RLController:
    def __init__(
        self,
        available_zones: List[str],
        log_period: Optional[int] = None,
        climate: Optional[str] = None,
        season: Optional[str] = None,
        selection_pickle: Optional[str] = None,
        device: torch.device = torch.device("cpu"),
        train_from_scratch: bool = False,
        policy_root_folder: str = "models/damper_control_policies",
    ):
        """
        Initialize the RL controller with the given parameters.

        Parameters
        ----------
        available_zones : List[str]
            List of available zones in the building
        log_period : Optional[int]
            Log period used for the RL controller selection
        climate : Optional[str]
            Climate zone used for the RL controller selection
        season : Optional[str]
            Season used for the RL controller selection
        selection_pickle : Optional[str]
            Path to the selection pickle file
        device : torch.device
            Device to run the RL controller
        train_from_scratch : bool
            Train the RL controller from scratch
        policy_root_folder : str
            Root folder for the policy
        """
        # Save the available zones
        self.available_zones = available_zones
        self.agents = dict()
        self.selected_policies = None
        self.agent_names = dict()

        # Load the policies selections
        if selection_pickle is not None and os.path.isfile(selection_pickle):
            try:
                print("Loading RL selections ...")
                with open(selection_pickle, "rb") as pklfile:
                    self.selected_policies = pickle.load(pklfile)[log_period][climate][season]
            except KeyError:
                print(f"Given selection pickle does not contain "
                      f"log_period={log_period}, climate={climate}, season={season} data")

        if self.selected_policies is None and not train_from_scratch:
            print("Using random RL selections ...")
            self.selected_policies = {
                zone: np.random.choice(os.listdir(policy_root_folder)) for zone in available_zones
            }

        if train_from_scratch:
            print("Initializing RL agents from scratch ...")

        # Load the policies
        for zone in available_zones:
            self.agents[zone] = PPO(
                state_dim=6,
                action_dim=1,
                lr_actor=0.003,
                lr_critic=0.0005,
                gamma=1,
                k_epochs=10,
                eps_clip=0.2,
                has_continuous_action_space=True,
                action_std_init=0.2,
                device=device,
                diverse_policies=list(),
                diverse_weight=0,
                diverse_increase=True
            )
            if not train_from_scratch:
                print(f"{policy_root_folder}/{self.selected_policies[zone]}")
                self.agents[zone].load(f"{policy_root_folder}/{self.selected_policies[zone]}")

    def get_actions(
        self,
        state: Dict[str, Any]
    ) -> Dict[str, float]:
        """
        Get the actions for the given state.

        Parameters
        ----------
        state : Dict[str, Any]
            Current state of the building

        Returns
        -------
        Dict[str, float]
            Dictionary containing the actions for each zone
        """
        actions = dict()
        for zone in self.available_zones:
            action = self.agents[zone].select_action(
                np.array([
                    state["outdoor temperature"],
                    state["outdoor solar"],
                    state["hour"],
                    state[f"{zone} humidity"],
                    state[f"{zone} thermostat temperature"],
                    state[f"{zone} occupancy"]
                ])
            )
            actions[zone] = 0.9 / (1 + np.exp(-action)) + 0.1
        return actions

    def get_action_means(
        self,
        state: Dict[str, Any]
    ) -> Dict[str, float]:
        """
        Get the action means for the given state.

        Parameters
        ----------
        state : Dict[str, Any]
            Current state of the building

        Returns
        -------
        Dict[str, float]
            action means for each zone
        """
        actions = dict()
        for zone in self.available_zones:
            action = self.agents[zone].get_mean(
                np.array([
                    state["outdoor temperature"],
                    state["outdoor solar"],
                    state["hour"],
                    state[f"{zone} humidity"],
                    state[f"{zone} thermostat temperature"],
                    state[f"{zone} occupancy"]
                ])
            )
            actions[zone] = 0.9 / (1 + np.exp(-action[0])) + 0.1
        return actions

    def evaluate_actions(
        self,
        state: Dict[str, Any],
        actions: Dict[str, float],
        next_state: Dict[str, Any],
    ) -> None:
        """
        Evaluate the actions for the given state.
        Parameters
        ----------
        state : Dict[str, Any]
            Current state of the building
        actions : Dict[str, float]
            Actions for each zone
        next_state : Dict[str, Any]
            Next state of the building
        """
        for zone in self.available_zones:
            self.agents[zone].learn_action_offline(
                np.array([
                    state["outdoor temperature"],
                    state["outdoor solar"],
                    state["hour"],
                    state[f"{zone} humidity"],
                    state[f"{zone} thermostat temperature"],
                    state[f"{zone} occupancy"]
                ]),
                actions[zone]
            )
            self.agents[zone].buffer.rewards.append(
                -next_state[f"{zone} Air System Sensible Cooling Rate"] +
                -next_state[f"{zone} Air System Sensible Heating Rate"]
            )
            self.agents[zone].buffer.is_terminals.append(False)

    def decay_action_std(
        self,
        action_std_decay_rate: float,
        min_action_std: float
    ) -> None:
        """
        Decay the action standard deviation for the agents.

        Parameters
        ----------
        action_std_decay_rate : float
            Decay rate for the action standard deviation
        min_action_std : float
            Minimum action standard deviation
        """
        for agent in self.agents.values():
            agent.decay_action_std(action_std_decay_rate, min_action_std)

    def update(self):
        """
        Update the agents.
        """
        for agent in self.agents.values():
            agent.update()


def extract_state(
    available_zones: List[str],
    state: Dict[str, Any]
) -> Dict[str, float]:
    """
    Extract the relevant state variables from the state dictionary.

    Parameters
    ----------
    available_zones : List[str]
        List of available zones in the building
    state : Dict[str, Any]
        Current state of the building

    Returns
    -------
    Dict[str, float]
        Dictionary containing the relevant state variables
    """
    cleaned_state = {"hour": state["time"].hour}
    for key in (
            "outdoor temperature",
            "outdoor wind",
            "outdoor solar",
            "total hvac electricity",
            "timestep",
            "day type"
    ):
        cleaned_state[key] = state[key]

    for zone in available_zones:
        for key in (
                f"{zone} thermostat heating setpoint",
                f"{zone} thermostat cooling setpoint",
                f"{zone} thermostat temperature",
                f"{zone} VAV inlet temp",
                f"{zone} humidity",
                f"{zone} total solar",
                f"{zone} Damper",
                f"{zone} thermostat temperature",
                f"{zone} Air System Sensible Cooling Rate",
                f"{zone} Air System Sensible Heating Rate"
        ):
            cleaned_state[key] = state[key]

        cleaned_state[f"{zone} predict load"] = abs(state[f"{zone} predict load"])
        cleaned_state[f"{zone} occupancy"] = state["occupancy"][zone]
    return cleaned_state


def icnn_inference(
    zone: str,
    state: Dict[str, float],
    prev_state: Dict[str, float],
    fg_model: ICNN,
    pretrained_models: Dict[str, List[OrderedDict[str, torch.Tensor]]],
    rff_functions: List[Callable]
) -> Tuple[Dict[str, np.ndarray], Dict[str, float]]:
    """
    Run ICNN inference for a specific zone and return the outputs and objectives.

    Parameters
    ----------
    zone : str
        Selected zone for ICNN inference
    state : Dict[str, float]
        Current state of the building
    prev_state : Dict[str, float]
        Previous state of the building
    fg_model : ICNN
        The ICNN model to be used for inference
    pretrained_models : Dict[str, List[OrderedDict[str, torch.Tensor]]]
        Pretrained ICNN model weights for F and G functions
    rff_functions : List[Callable]
        List of Random Fourier Features (RFF) functions for G estimation

    Returns
    -------
    Tuple[Dict[str, np.ndarray], Dict[str, float]]
        Tuple containing the following two dictionaries:
            1. ICNN outputs for each function (F and G)
            2. Objectives for each function (F and G)
    """
    icnn_outputs = dict()
    objectives = dict()

    # load ICNN files and generate outputs from ICNN based on the states
    predictor_input = np.array([
        prev_state["outdoor temperature"],
        prev_state["outdoor wind"],
        prev_state["outdoor solar"],
        prev_state["hour"],
        prev_state[f"{zone} occupancy"],
        prev_state[f"{zone} thermostat heating setpoint"],
        prev_state[f"{zone} thermostat cooling setpoint"],
        prev_state[f"{zone} thermostat temperature"],
        prev_state[f"{zone} VAV inlet temp"],
        abs(prev_state[f"{zone} predict load"]),
        prev_state[f"{zone} total solar"],
        prev_state[f"{zone} Damper"]
    ]).astype("float32")

    for key in "fg":
        current_icnn_outputs = list()
        # ICNN inference
        for ii in range(len(pretrained_models[key])):
            # Load the ICNN model
            fg_model.load_state_dict(pretrained_models[key][ii])
            fg_model.eval()

            # Run ICNN and scale the output
            fg_model_output = fg_model(torch.Tensor(predictor_input).to("cpu"))[0].item()
            scale_min, scale_max = pretrained_models[f"{key}_scale"][ii]
            current_icnn_outputs.append(fg_model_output * (scale_max - scale_min) + scale_min)
        current_icnn_outputs = np.array(current_icnn_outputs)

        # Calculate the objective based on the ICNN output
        if key == 'f':
            objective = state[f"{zone} thermostat temperature"]
            icnn_outputs[key] = current_icnn_outputs
        elif key == 'g':
            objective = abs(state[f"{zone} Air System Sensible Cooling Rate"] +
                            state[f"{zone} Air System Sensible Heating Rate"])
            rff_evals = np.array([f(predictor_input) for f in rff_functions])
            icnn_outputs[key] = np.hstack([current_icnn_outputs, rff_evals])
        else:
            raise ValueError(f"Unknown ICNN key {key}")

        objectives[key] = objective
    return icnn_outputs, objectives


def set_extra_states(
    available_zones: List[str]
) -> Dict[Tuple[str, str], str]:
    """
    Set up the extra states for the EnergyPlus simulation.

    Parameters
    ----------
    available_zones : List[str]
        List of available zones in the building

    Returns
    -------
    Dict[Tuple[str, str], str]
        Dictionary containing the extra states for the EnergyPlus simulation
    """
    # Set up zone level readings
    eplus_extra_states = dict()
    for zone in available_zones:
        eplus_extra_states[('Zone Air Terminal VAV Damper Position', f"{zone} VAV Box Component")] = f"{zone} Damper"
        eplus_extra_states[("System Node Temperature", f"{zone} VAV Box Damper Node")] = f"{zone} VAV inlet temp"
        eplus_extra_states[("System Node Temperature", f"{zone} VAV Box Outlet Node")] = f"{zone} VAV outlet temp"
        for stat_name, save_name in [
            ('Zone Thermostat Heating Setpoint Temperature', f"{zone} thermostat heating setpoint"),
            ('Zone Thermostat Cooling Setpoint Temperature', f"{zone} thermostat cooling setpoint"),
            ('Zone Thermostat Air Temperature', f"{zone} thermostat temperature"),
            ("Zone Windows Total Transmitted Solar Radiation Rate", f"{zone} total solar"),
            ("Zone Predicted Sensible Load to Setpoint Heat Transfer Rate", f"{zone} predict load"),
            ("Zone Air Relative Humidity", f"{zone} humidity"),
        ]:
            eplus_extra_states[(stat_name, zone)] = save_name
        for name in [
            "Zone Air System Sensible Heating Energy", "Zone Air System Sensible Heating Rate",
            "Zone Air System Sensible Cooling Energy", "Zone Air System Sensible Cooling Rate",
        ]:
            eplus_extra_states[(name, zone)] = name.replace("Zone", zone)

    # Set up building level readings
    eplus_extra_states[('Site Outdoor Air Drybulb Temperature', 'Environment')] = "outdoor temperature"
    eplus_extra_states[('Site Outdoor Air Humidity Ratio', 'Environment')] = "outdoor humidity"
    eplus_extra_states[('Site Wind Speed', 'Environment')] = "outdoor wind"
    eplus_extra_states[('Site Direct Solar Radiation Rate per Area', 'Environment')] = "outdoor solar"
    eplus_extra_states[('Site Solar Hour Angle', 'Environment')] = "solar angle"
    eplus_extra_states[('Site Precipitation Depth', 'Environment')] = "precipitation"
    eplus_extra_states[('Site Day Type Index', 'Environment')] = "day type"
    eplus_extra_states[('Facility Total HVAC Electric Demand Power', 'Whole Building')] = "total hvac electricity"

    return eplus_extra_states


def set_cli() -> argparse.ArgumentParser:
    """
    Set up the command line interface for the experiment.

    Returns
    -------
    argparse.ArgumentParser
        Argument parser for the experiment
    """
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--experiment_duration',
        help="Duration of the experiment in days",
        type=int,
        default=21,
    )
    parser.add_argument(
        '--experiment_interval',
        help="How many steps in an hour",
        type=int,
        default=4,
    )
    parser.add_argument(
        '--year',
        help="Experiment start year",
        type=int,
        default=2000,
    )
    parser.add_argument(
        '--seed',
        help="random seed",
        type=int,
        default=3,
    )
    parser.add_argument(
        '--month',
        help="Experiment start month",
        type=int,
        default=1,
    )
    parser.add_argument(
        '--day',
        help="Experiment start day",
        type=int,
        default=15,
    )
    parser.add_argument(
        '--log_days',
        help="How many days before the experiment used to collect log data",
        type=int,
        default=14,
    )
    parser.add_argument(
        '--icnn_random',
        help="Experiment random ICNNs",
        type=int,
        default=0,
    )
    parser.add_argument(
        '--experiment_building',
        help="Building name to run the experiment",
        type=str,
        default="target_building",
    )
    parser.add_argument(
        '--experiment_season',
        help="Season to run experiment, will overwrite some of the parameters",
        type=str,
        choices=['summer', 'winter'],
        default=None,
    )
    parser.add_argument(
        '--experiment_climate',
        help="Climate zone to run experiment, will overwrite some of the parameters",
        type=str,
        choices=["0a", "0b", "1a", "1b", "2a", "2b", "3a", "3b", "3c",
                 "4a", "4b", "4c", "5a", "5b", "5c", "6a", "6b", "7", "8"],
        default="5b",
    )
    parser.add_argument(
        '--get_log',
        help="Collect log data or run experiment",
        action="store_true",
    )
    parser.add_argument(
        '--logfile_name',
        help="Name of the log file",
        type=str,
        default="Jan01-2000",
    )
    parser.add_argument(
        '--savedir_suffix',
        help="Suffix to distinguish duplicate runs",
        type=str,
        default="",
    )
    parser.add_argument(
        '--dir_prefix',
        help="Need for Ray",
        type=str,
        default=".",
    )
    parser.add_argument(
        '--w',
        help="W_f parameter",
        type=float,
        default=1.5,
    )
    parser.add_argument(
        '--rl_lambda',
        help="RL lambda parameter",
        type=int,
        default=50,
    )
    parser.add_argument(
        '--rl_train_episodes',
        help="Number of episodes to train RL baselines",
        type=int,
        default=1000,
    )
    parser.add_argument(
        "--train_CRL",
        action="store_true",
        help="Code in training ConstraintRL mode",
    )
    parser.add_argument(
        "--remove_optimization",
        action="store_true",
        help="Remove the optimization after transfer RL baselines",
    )
    parser.add_argument(
        "--update_optimization",
        action="store_true",
        help="Re-identify the optimization after transfer RL baselines",
    )
    parser.add_argument(
        '--icnn_size',
        help="number of icnns to select",
        type=int,
        default=60,
    )
    parser.add_argument(
        "--dry_run",
        action="store_true",
        help="The model runs locally and the local dataset already "
             "exists and does not need to be downloaded from S3",
    )
    parser.add_argument(
        "--lse",
        action="store_true",
        help="The model run lse baseline",
    )
    parser.add_argument(
        "--cls",
        action="store_true",
        help="The model run cls baseline",
    )
    parser.add_argument(
        "--coef_ls",
        help="number of icnns to select",
        type=float,
        default=20,
    )
    parser.add_argument(
        "--verbose",
        help="The print level",
        type=int,
        default=0,
    )
    parser.add_argument(
        "--multiprocessing",
        action="store_true",
        help="The model use multiprocessing to speed up",
    )
    return parser


def set_model(
    available_zones: List[str],
    args: argparse.Namespace
) -> Model:
    """
    Set up the EnergyPlus model for the experiment.

    Parameters
    ----------
    available_zones : List[str]
        List of available zones in the building
    args : argparse.Namespace
        Arguments for the experiment

    Returns
    -------
    Model
        EnergyPlus model for the experiment
    """
    # Set up zone level readings
    eplus_extra_states = set_extra_states(available_zones)
    random_suffix = ''.join(random.choices(string.ascii_letters + string.digits, k=7))
    log_folder = f"./results/{datetime.now().strftime('%Y%m%d_%H%M')}_{random_suffix}"
    os.makedirs(f"{log_folder}/eplus_output", exist_ok=True)

    model = Model(
        idf_file_name=f"./eplus_files/{args.experiment_building}.idf",
        weather_file=f"./eplus_files/weather/zone_{args.experiment_climate}.epw",
        eplus_naming_dict=eplus_extra_states,
        tmp_idf_path=f"{log_folder}/eplus_output/"
    )

    # Add sensors to the IDF file, so we can retrieve them later
    for key, _ in eplus_extra_states.items():
        model.add_configuration(
            idf_header_name="Output:Variable",
            values={
                "Key Value": key[1],
                "Variable Name": key[0],
                "Reporting Frequency": "Timestep"
            }
        )

    # Add VAV boxes and control to the IDF file
    if not args.get_log:
        for zone in available_zones:
            model.add_configuration(
                idf_header_name="Schedule:Constant",
                values={
                    "Name": f"{zone} VAV Customized Schedule",
                    "Schedule Type Limits Name": "Fraction",
                    "Hourly Value": 0
                }
            )
            model.edit_configuration(
                idf_header_name="AirTerminal:SingleDuct:VAV:Reheat",
                identifier={"Name": f"{zone} VAV Box Component"},
                update_values={
                    "Zone Minimum Air Flow Input Method": "Scheduled",
                    "Minimum Air Flow Fraction Schedule Name": f"{zone} VAV Customized Schedule"
                }
            )

    # Simulate runDuration days, starting from runMonth runDate in runYear
    model.set_runperiod(
        start_year=args.year,
        start_month=args.month,
        start_day=args.day,
        days=args.experiment_duration,
    )
    # Each hour contains runInterval time steps (runInterval = 4 --> 15 min per time step)
    model.set_timestep(args.experiment_interval)

    return model