MIP_DQN.py

# ------------------------------------------------------------------------
# Optimal-Energy-System-Scheduling-Combining-Mixed-Integer-Programming-and-Deep-Reinforcement-Learning
# MIP-DQN algorithm developed by
# Hou Shengren, TU Delft, h.shengren@tudelft.nl
# Pedro, TU Delft, p.p.vergara.barrios@tudeflt.nl
# ------------------------------------------------------------------------
import os
import pickle
import tempfile
from copy import deepcopy
from pathlib import Path

import numpy as np
import numpy.random as rd
import torch
import torch.nn as nn
import torch.onnx

try:
    import pyomo.environ as pyo
except ImportError:  # pragma: no cover - optional dependency
    pyo = None

try:
    from omlt import OmltBlock
    from omlt.io.onnx import load_onnx_neural_network_with_bounds, write_onnx_model_with_bounds
    from omlt.neuralnet import ReluBigMFormulation
except ImportError:  # pragma: no cover - optional dependency
    OmltBlock = None
    ReluBigMFormulation = None
    load_onnx_neural_network_with_bounds = None
    write_onnx_model_with_bounds = None

try:
    import wandb
except ImportError:  # pragma: no cover - optional dependency
    wandb = None

from random_generator_battery import ESSEnv


def _env_flag(name, default):
    value = os.getenv(name)
    if value is None:
        return default
    return value.lower() in {"1", "true", "yes", "on"}


def _env_int(name, default):
    value = os.getenv(name)
    return int(value) if value is not None else default


def _env_float(name, default):
    value = os.getenv(name)
    return float(value) if value is not None else default


def _env_int_list(name, default):
    value = os.getenv(name)
    if not value:
        return default
    return [int(item.strip()) for item in value.split(",") if item.strip()]


class ReplayBuffer:
    def __init__(self, max_len, state_dim, action_dim, gpu_id=0):
        self.now_len = 0
        self.next_idx = 0
        self.if_full = False
        self.max_len = max_len
        self.data_type = torch.float32
        self.action_dim = action_dim
        self.device = torch.device(
            f"cuda:{gpu_id}" if (torch.cuda.is_available() and (gpu_id >= 0)) else "cpu"
        )

        other_dim = 1 + 1 + self.action_dim
        self.buf_other = torch.empty(size=(max_len, other_dim), dtype=self.data_type, device=self.device)

        if isinstance(state_dim, int):
            self.buf_state = torch.empty((max_len, state_dim), dtype=torch.float32, device=self.device)
        elif isinstance(state_dim, tuple):
            self.buf_state = torch.empty((max_len, *state_dim), dtype=torch.uint8, device=self.device)
        else:
            raise ValueError("state_dim")

    def extend_buffer(self, state, other):
        size = len(other)
        next_idx = self.next_idx + size

        if next_idx > self.max_len:
            self.buf_state[self.next_idx : self.max_len] = state[: self.max_len - self.next_idx]
            self.buf_other[self.next_idx : self.max_len] = other[: self.max_len - self.next_idx]
            self.if_full = True

            next_idx = next_idx - self.max_len
            self.buf_state[0:next_idx] = state[-next_idx:]
            self.buf_other[0:next_idx] = other[-next_idx:]
        else:
            self.buf_state[self.next_idx:next_idx] = state
            self.buf_other[self.next_idx:next_idx] = other
        self.next_idx = next_idx

    def sample_batch(self, batch_size):
        if self.now_len <= 1:
            raise RuntimeError("ReplayBuffer does not contain enough transitions to sample a batch.")
        indices = rd.randint(self.now_len - 1, size=batch_size)
        r_m_a = self.buf_other[indices]
        return (
            r_m_a[:, 0:1],
            r_m_a[:, 1:2],
            r_m_a[:, 2:],
            self.buf_state[indices],
            self.buf_state[indices + 1],
        )

    def update_now_len(self):
        self.now_len = self.max_len if self.if_full else self.next_idx


class Arguments:
    def __init__(self, agent=None, env=None):
        self.agent = agent
        self.env = env
        self.cwd = None
        self.if_remove = False
        self.visible_gpu = os.getenv("MIP_DQN_VISIBLE_GPU", os.getenv("CUDA_VISIBLE_DEVICES", "0"))
        self.worker_num = 2
        self.num_threads = _env_int("MIP_DQN_NUM_THREADS", 8)

        self.num_episode = _env_int("MIP_DQN_NUM_EPISODES", 3000)
        self.gamma = _env_float("MIP_DQN_GAMMA", 0.995)
        self.learning_rate = _env_float("MIP_DQN_LEARNING_RATE", 1e-4)
        self.soft_update_tau = _env_float("MIP_DQN_SOFT_UPDATE_TAU", 1e-2)

        self.net_dim = _env_int("MIP_DQN_NET_DIM", 64)
        self.batch_size = _env_int("MIP_DQN_BATCH_SIZE", 256)
        self.repeat_times = _env_int("MIP_DQN_REPEAT_TIMES", 2**3)
        self.target_step = _env_int("MIP_DQN_TARGET_STEP", 1000)
        self.max_memo = _env_int("MIP_DQN_MAX_MEMO", 50000)
        self.initial_buffer_size = _env_int("MIP_DQN_INITIAL_BUFFER_SIZE", 10000)
        self.explorate_decay = _env_float("MIP_DQN_EXPLORATE_DECAY", 0.99)
        self.explorate_min = _env_float("MIP_DQN_EXPLORATE_MIN", 0.3)

        self.random_seed_list = _env_int_list("MIP_DQN_RANDOM_SEEDS", [1234, 2234, 3234, 4234, 5234])
        self.run_name = os.getenv("MIP_DQN_RUN_NAME", "MIP_DQN_experiments")

        self.train = _env_flag("MIP_DQN_TRAIN", True)
        self.save_network = _env_flag("MIP_DQN_SAVE_NETWORK", True)
        self.update_training_data = _env_flag("MIP_DQN_SAVE_RECORDS", True)
        self.enable_wandb = _env_flag("MIP_DQN_ENABLE_WANDB", wandb is not None)
        self.use_actor_mip = _env_flag("MIP_DQN_USE_ACTOR_MIP", True)
        self.actor_mip_projection_radius = _env_float("MIP_DQN_PROJECTION_RADIUS", 0.75)
        self.if_per_or_gae = False

    def init_before_training(self, if_main):
        if self.cwd is None:
            agent_name = self.agent.__class__.__name__
            self.cwd = f"./{agent_name}/{self.run_name}/seed_{self.random_seed}"

        if if_main:
            import shutil

            if self.if_remove is None:
                self.if_remove = bool(input(f"| PRESS 'y' to REMOVE: {self.cwd}? ") == "y")
            elif self.if_remove:
                shutil.rmtree(self.cwd, ignore_errors=True)
                print(f"| Remove cwd: {self.cwd}")
            os.makedirs(self.cwd, exist_ok=True)

        np.random.seed(self.random_seed)
        torch.manual_seed(self.random_seed)
        torch.set_num_threads(self.num_threads)
        torch.set_default_dtype(torch.float32)
        os.environ["CUDA_VISIBLE_DEVICES"] = str(self.visible_gpu)


class Actor(nn.Module):
    def __init__(self, mid_dim, state_dim, action_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim, mid_dim),
            nn.ReLU(),
            nn.Linear(mid_dim, mid_dim),
            nn.ReLU(),
            nn.Linear(mid_dim, mid_dim),
            nn.ReLU(),
            nn.Linear(mid_dim, action_dim),
        )

    def forward(self, state):
        return self.net(state).tanh()

    def get_action(self, state, action_std):
        action = self.net(state).tanh()
        noise = (torch.randn_like(action) * action_std).clamp(-0.5, 0.5)
        return (action + noise).clamp(-1.0, 1.0)


class CriticQ(nn.Module):
    def __init__(self, mid_dim, state_dim, action_dim):
        super().__init__()
        self.net_head = nn.Sequential(
            nn.Linear(state_dim + action_dim, mid_dim),
            nn.ReLU(),
            nn.Linear(mid_dim, mid_dim),
            nn.ReLU(),
        )
        self.net_q1 = nn.Sequential(nn.Linear(mid_dim, mid_dim), nn.ReLU(), nn.Linear(mid_dim, 1))
        self.net_q2 = nn.Sequential(nn.Linear(mid_dim, mid_dim), nn.ReLU(), nn.Linear(mid_dim, 1))

    def forward(self, value):
        mid = self.net_head(value)
        return self.net_q1(mid)

    def get_q1_q2(self, value):
        mid = self.net_head(value)
        return self.net_q1(mid), self.net_q2(mid)


class AgentBase:
    def __init__(self):
        self.state = None
        self.device = None
        self.action_dim = None
        self.if_off_policy = None
        self.explore_noise = None
        self.trajectory_list = None
        self.explore_rate = 1.0

        self.criterion = torch.nn.SmoothL1Loss()

    def init(self, net_dim, state_dim, action_dim, learning_rate=1e-4, _if_per_or_gae=False, gpu_id=0):
        self.device = torch.device(
            f"cuda:{gpu_id}" if (torch.cuda.is_available() and (gpu_id >= 0)) else "cpu"
        )
        self.action_dim = action_dim

        self.cri = self.ClassCri(net_dim, state_dim, action_dim).to(self.device)
        self.act = self.ClassAct(net_dim, state_dim, action_dim).to(self.device) if self.ClassAct else self.cri
        self.cri_target = deepcopy(self.cri) if self.if_use_cri_target else self.cri
        self.act_target = deepcopy(self.act) if self.if_use_act_target else self.act

        self.cri_optim = torch.optim.Adam(self.cri.parameters(), learning_rate)
        self.act_optim = torch.optim.Adam(self.act.parameters(), learning_rate) if self.ClassAct else self.cri
        del self.ClassCri, self.ClassAct

    def get_raw_action(self, state, use_exploration=True):
        states = torch.as_tensor(np.asarray(state, dtype=np.float32)[None, :], dtype=torch.float32, device=self.device)
        action = self.act(states)[0]
        if use_exploration and rd.rand() < self.explore_rate:
            action = (action + torch.randn_like(action) * self.explore_noise).clamp(-1, 1)
        return action.detach().cpu().numpy()

    def select_action(self, state, actor_mip=None, use_exploration=True):
        raw_action = self.get_raw_action(state, use_exploration=use_exploration)
        if actor_mip is None:
            return raw_action
        return actor_mip.project_action(state, raw_action)

    def explore_env(self, env, target_step, actor_mip=None):
        trajectory = []
        state = self.state

        for _ in range(target_step):
            action = self.select_action(state, actor_mip=actor_mip, use_exploration=True)
            current_state, next_state, reward, done = env.step(action)
            trajectory.append((current_state, (reward, done, *action)))
            state = env.reset() if done else next_state
        self.state = state
        return trajectory

    @staticmethod
    def optim_update(optimizer, objective):
        optimizer.zero_grad()
        objective.backward()
        optimizer.step()

    @staticmethod
    def soft_update(target_net, current_net, tau):
        for tar, cur in zip(target_net.parameters(), current_net.parameters()):
            tar.data.copy_(cur.data * tau + tar.data * (1.0 - tau))

    def save_or_load_agent(self, cwd, if_save):
        def load_torch_file(model_or_optim, path):
            state_dict = torch.load(path, map_location=lambda storage, loc: storage)
            model_or_optim.load_state_dict(state_dict)

        name_obj_list = [
            ("actor", self.act),
            ("act_target", self.act_target),
            ("act_optim", self.act_optim),
            ("critic", self.cri),
            ("cri_target", self.cri_target),
            ("cri_optim", self.cri_optim),
        ]
        name_obj_list = [(name, obj) for name, obj in name_obj_list if obj is not None]
        if if_save:
            for name, obj in name_obj_list:
                save_path = f"{cwd}/{name}.pth"
                torch.save(obj.state_dict(), save_path)
        else:
            for name, obj in name_obj_list:
                save_path = f"{cwd}/{name}.pth"
                if os.path.isfile(save_path):
                    load_torch_file(obj, save_path)

    def _update_exploration_rate(self, explorate_decay, explore_rate_min):
        self.explore_rate = max(self.explore_rate * explorate_decay, explore_rate_min)


class AgentMIPDQN(AgentBase):
    def __init__(self):
        super().__init__()
        self.explore_noise = 0.5
        self.policy_noise = 0.2
        self.update_freq = 2
        self.if_use_cri_target = self.if_use_act_target = True
        self.ClassCri = CriticQ
        self.ClassAct = Actor

    def update_net(self, buffer, batch_size, repeat_times, soft_update_tau):
        buffer.update_now_len()
        obj_critic = obj_actor = None
        update_steps = int(buffer.now_len / batch_size * repeat_times)
        for update_c in range(update_steps):
            obj_critic, state = self.get_obj_critic(buffer, batch_size)
            self.optim_update(self.cri_optim, obj_critic)

            action_pg = self.act(state)
            obj_actor = -self.cri_target(torch.cat((state, action_pg), dim=-1)).mean()
            self.optim_update(self.act_optim, obj_actor)
            if update_c % self.update_freq == 0:
                self.soft_update(self.cri_target, self.cri, soft_update_tau)
                self.soft_update(self.act_target, self.act, soft_update_tau)
        return obj_critic.item() / 2, obj_actor.item()

    def get_obj_critic(self, buffer, batch_size):
        with torch.no_grad():
            reward, mask, action, state, next_s = buffer.sample_batch(batch_size)
            next_a = self.act_target.get_action(next_s, self.policy_noise)
            next_q = torch.min(*self.cri_target.get_q1_q2(torch.cat((next_s, next_a), dim=-1)))
            q_label = reward + mask * next_q

        q1, q2 = self.cri.get_q1_q2(torch.cat((state, action), dim=-1))
        obj_critic = self.criterion(q1, q_label) + self.criterion(q2, q_label)
        return obj_critic, state


def update_buffer(buffer, trajectory, gamma):
    ten_state = torch.as_tensor(np.asarray([item[0] for item in trajectory], dtype=np.float32), dtype=torch.float32)
    ary_other = torch.as_tensor(np.asarray([item[1] for item in trajectory], dtype=np.float32), dtype=torch.float32)
    ary_other[:, 1] = (1.0 - ary_other[:, 1]) * gamma

    buffer.extend_buffer(ten_state, ary_other)

    steps = ten_state.shape[0]
    r_exp = ary_other[:, 0].mean()
    return steps, r_exp


def get_episode_return(env, agent, actor_mip=None):
    episode_return = 0.0
    episode_unbalance = 0.0
    episode_operation_cost = 0.0
    state = env.reset()
    for _ in range(env.episode_length):
        action = agent.select_action(state, actor_mip=actor_mip, use_exploration=False)
        _, next_state, reward, done = env.step(action)
        state = next_state
        episode_return += reward
        episode_unbalance += env.real_unbalance
        episode_operation_cost += env.operation_cost
        if done:
            break
    return episode_return, episode_unbalance, episode_operation_cost


class Actor_MIP:
    SCALE_NAMES = (
        "battery_power",
        "dg1_ramping",
        "dg2_ramping",
        "dg3_ramping",
        "net_load_scale",
        "dg1_output_scale",
        "dg2_output_scale",
        "dg3_output_scale",
    )

    def __init__(
        self,
        scaled_parameters,
        batch_size,
        net,
        state_dim,
        action_dim,
        env,
        constrain_on=False,
        projection_radius=1.0,
        solver_name="gurobi",
    ):
        self.batch_size = batch_size
        self.net = net
        self.state_dim = state_dim
        self.action_dim = action_dim
        self.env = env
        self.constrain_on = constrain_on
        self.projection_radius = projection_radius
        self.solver_name = solver_name
        self.scaled_parameters = self._normalize_scaled_parameters(scaled_parameters)

        self.state_names = tuple(getattr(env, "state_names", ()))
        self.action_names = tuple(getattr(env, "action_names", ()))
        if len(self.state_names) != self.state_dim or len(self.action_names) != self.action_dim:
            raise ValueError("Environment state/action ordering is inconsistent with the network dimensions.")

        self.state_indices = {name: offset for offset, name in enumerate(self.state_names)}
        self.action_indices = {name: self.state_dim + offset for offset, name in enumerate(self.action_names)}

        self.cache_stats = {"onnx_exports": 0, "model_builds": 0}
        self._solver = None
        self._model = None
        self._network_signature = None
        self._cached_export_path = None

    def _normalize_scaled_parameters(self, scaled_parameters):
        if isinstance(scaled_parameters, dict):
            missing = [name for name in self.SCALE_NAMES if name not in scaled_parameters]
            if missing:
                raise ValueError(f"scaled_parameters is missing keys: {missing}")
            return {name: float(scaled_parameters[name]) for name in self.SCALE_NAMES}

        if len(scaled_parameters) != len(self.SCALE_NAMES):
            raise ValueError("scaled_parameters must contain 8 elements.")
        return {name: float(value) for name, value in zip(self.SCALE_NAMES, scaled_parameters)}

    def _current_network_signature(self):
        return id(self.net), tuple(parameter._version for parameter in self.net.parameters())

    def sync_network(self, net=None):
        if net is not None:
            self.net = net
        signature = self._current_network_signature()
        if signature != self._network_signature:
            self._network_signature = signature
            self._invalidate_cache()

    def _invalidate_cache(self):
        self._model = None
        if self._cached_export_path and os.path.exists(self._cached_export_path):
            os.remove(self._cached_export_path)
        self._cached_export_path = None

    def get_input_bounds(self, input_batch_state, input_batch_action=None):
        batch_state = np.atleast_2d(np.asarray(input_batch_state, dtype=np.float32))
        if input_batch_action is None:
            batch_action = np.zeros((batch_state.shape[0], self.action_dim), dtype=np.float32)
        else:
            batch_action = np.atleast_2d(np.asarray(input_batch_action, dtype=np.float32))

        bounds = []
        state_low = np.asarray(self.env.state_space.low, dtype=np.float32)
        state_high = np.asarray(self.env.state_space.high, dtype=np.float32)
        for index in range(batch_state.shape[0]):
            bound = {}
            for state_idx in range(self.state_dim):
                bound[state_idx] = (float(state_low[state_idx]), float(state_high[state_idx]))
            action_lower, action_upper = self._build_projection_bounds(batch_action[index])
            for action_idx in range(self.action_dim):
                bound[self.state_dim + action_idx] = (
                    float(action_lower[action_idx]),
                    float(action_upper[action_idx]),
                )
            bounds.append(bound)
        return bounds

    def _require_optional_dependencies(self):
        if (
            pyo is None
            or OmltBlock is None
            or ReluBigMFormulation is None
            or load_onnx_neural_network_with_bounds is None
            or write_onnx_model_with_bounds is None
        ):
            raise RuntimeError(
                "Actor_MIP requires pyomo, omlt, and their ONNX helpers to be installed."
            )

    def _ensure_solver(self):
        self._require_optional_dependencies()
        if self._solver is None:
            solver = pyo.SolverFactory(self.solver_name)
            if solver is None or not solver.available(exception_flag=False):
                raise RuntimeError(
                    f"Actor_MIP requires the '{self.solver_name}' solver to be available in the environment."
                )
            self._solver = solver
        return self._solver

    def _export_network_to_onnx(self):
        dummy_input = torch.zeros((1, self.state_dim + self.action_dim), dtype=torch.float32)
        model = deepcopy(self.net).to("cpu").eval()
        with tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) as handle:
            export_path = handle.name

        try:
            torch.onnx.export(
                model,
                dummy_input,
                export_path,
                input_names=["state_action"],
                output_names=["Q_value"],
                dynamic_axes={"state_action": {0: "batch_size"}, "Q_value": {0: "batch_size"}},
            )
        except Exception as exc:  # pragma: no cover - depends on optional runtime stack
            if os.path.exists(export_path):
                os.remove(export_path)
            raise RuntimeError(
                "Actor_MIP requires the ONNX export stack used by torch.onnx. "
                "Install the missing ONNX dependencies before enabling Actor_MIP."
            ) from exc
        self.cache_stats["onnx_exports"] += 1
        return export_path

    def _generic_input_bounds(self):
        bounds = {}
        state_low = np.asarray(self.env.state_space.low, dtype=np.float32)
        state_high = np.asarray(self.env.state_space.high, dtype=np.float32)
        for index in range(self.state_dim):
            bounds[index] = (float(state_low[index]), float(state_high[index]))
        for index in range(self.action_dim):
            bounds[self.state_dim + index] = (-1.0, 1.0)
        return bounds

    def _ensure_model(self):
        self._require_optional_dependencies()
        self.sync_network()
        if self._model is not None:
            return self._model

        export_path = self._export_network_to_onnx()
        write_onnx_model_with_bounds(export_path, None, self._generic_input_bounds())
        network_definition = load_onnx_neural_network_with_bounds(export_path)
        formulation = ReluBigMFormulation(network_definition)
        model = pyo.ConcreteModel()
        model.nn = OmltBlock()
        model.nn.build_formulation(formulation)

        if self.constrain_on:
            self._add_projection_constraints(model)

        model.obj = pyo.Objective(expr=model.nn.outputs[0], sense=pyo.maximize)
        self._model = model
        self._cached_export_path = export_path
        self.cache_stats["model_builds"] += 1
        return self._model

    def _state_input(self, model, name):
        return model.nn.inputs[self.state_indices[name]]

    def _action_input(self, model, name):
        return model.nn.inputs[self.action_indices[name]]

    def _add_projection_constraints(self, model):
        battery_action = self._action_input(model, "battery")
        dg1_action = self._action_input(model, "dg1")
        dg2_action = self._action_input(model, "dg2")
        dg3_action = self._action_input(model, "dg3")

        dg1_output = self._state_input(model, "dg1_output")
        dg2_output = self._state_input(model, "dg2_output")
        dg3_output = self._state_input(model, "dg3_output")
        net_load = self._state_input(model, "net_load")

        scales = self.scaled_parameters
        projected_balance = (
            -battery_action * scales["battery_power"]
            + (dg1_action * scales["dg1_ramping"] + dg1_output * scales["dg1_output_scale"])
            + (dg2_action * scales["dg2_ramping"] + dg2_output * scales["dg2_output_scale"])
            + (dg3_action * scales["dg3_ramping"] + dg3_output * scales["dg3_output_scale"])
        )
        lower_bound = net_load * scales["net_load_scale"] - self.env.grid.exchange_ability
        upper_bound = net_load * scales["net_load_scale"] + self.env.grid.exchange_ability

        model.power_balance_lower = pyo.Constraint(expr=projected_balance >= lower_bound)
        model.power_balance_upper = pyo.Constraint(expr=projected_balance <= upper_bound)

    def _build_projection_bounds(self, actor_action):
        actor_action = np.asarray(actor_action, dtype=np.float32).reshape(self.action_dim)
        if self.projection_radius >= 1.0:
            lower = np.full(self.action_dim, -1.0, dtype=np.float32)
            upper = np.full(self.action_dim, 1.0, dtype=np.float32)
        else:
            lower = np.clip(actor_action - self.projection_radius, -1.0, 1.0)
            upper = np.clip(actor_action + self.projection_radius, -1.0, 1.0)
        return lower, upper

    def _bind_state(self, model, state):
        state = np.asarray(state, dtype=np.float32).reshape(self.state_dim)
        for offset, value in enumerate(state):
            model.nn.inputs[offset].fix(float(value))

    def _bind_action_search_region(self, model, actor_action):
        actor_action = np.asarray(actor_action, dtype=np.float32).reshape(self.action_dim)
        lower, upper = self._build_projection_bounds(actor_action)
        for offset, value in enumerate(actor_action):
            variable = model.nn.inputs[self.state_dim + offset]
            variable.unfix()
            variable.setlb(float(lower[offset]))
            variable.setub(float(upper[offset]))
            variable.set_value(float(value))

    def _solve_current_model(self, state, actor_action):
        model = self._ensure_model()
        self._bind_state(model, state)
        self._bind_action_search_region(model, actor_action)

        results = self._ensure_solver().solve(model, tee=False)
        status = results.solver.status
        termination = results.solver.termination_condition
        if status != pyo.SolverStatus.ok or termination != pyo.TerminationCondition.optimal:
            raise RuntimeError(
                f"Actor_MIP solve failed with status={status}, termination={termination}, "
                f"state={np.asarray(state).tolist()}, actor_action={np.asarray(actor_action).tolist()}."
            )
        return model

    def project_action(self, state, actor_action):
        model = self._solve_current_model(state, actor_action)
        action_values = []
        for offset in range(self.action_dim):
            variable = model.nn.inputs[self.state_dim + offset]
            action_values.append(float(pyo.value(variable)))
        return np.asarray(action_values, dtype=np.float32)

    def predict_best_action(self, state, actor_action=None):
        if actor_action is None:
            actor_action = np.zeros(self.action_dim, dtype=np.float32)
        return self.project_action(state, actor_action)


class NullWandbRun:
    def log(self, *_args, **_kwargs):
        return None

    def finish(self):
        return None


def init_wandb_run(args):
    if not args.enable_wandb:
        return NullWandbRun()
    if wandb is None:
        print("wandb is not installed, proceeding without experiment logging.")
        return NullWandbRun()

    settings = wandb.Settings(start_method="fork")
    run = wandb.init(
        project="MIP_DQN_experiments",
        name=f"{args.run_name}_seed_{args.random_seed}",
        settings=settings,
    )
    run.config.update({"epochs": args.num_episode, "batch_size": args.batch_size}, allow_val_change=True)
    wandb.define_metric("custom_step")
    return run


def build_actor_mip(agent, env, args):
    if not args.use_actor_mip:
        return None
    return Actor_MIP(
        scaled_parameters=env.get_actor_mip_config()["scaled_parameters"],
        batch_size=1,
        net=agent.cri_target,
        state_dim=env.state_space.shape[0],
        action_dim=env.action_space.shape[0],
        env=env,
        constrain_on=True,
        projection_radius=args.actor_mip_projection_radius,
    )


def collect_trajectory(agent, env, buffer, args, actor_mip):
    with torch.no_grad():
        trajectory = agent.explore_env(env, args.target_step, actor_mip=actor_mip)
        steps, r_exp = update_buffer(buffer, trajectory, args.gamma)
        buffer.update_now_len()
    return steps, r_exp


def save_training_outputs(cwd, args, agent, loss_record, reward_record):
    cwd_path = Path(cwd)
    if args.update_training_data:
        with (cwd_path / "loss_data.pkl").open("wb") as loss_file:
            pickle.dump(loss_record, loss_file)
        with (cwd_path / "reward_data.pkl").open("wb") as reward_file:
            pickle.dump(reward_record, reward_file)

    if args.save_network:
        torch.save(agent.act.state_dict(), cwd_path / "actor.pth")
        torch.save(agent.cri.state_dict(), cwd_path / "critic.pth")


def train_one_seed(seed, args):
    args.random_seed = seed
    args.agent = AgentMIPDQN()
    args.env = ESSEnv()
    args.cwd = None
    args.init_before_training(if_main=True)

    agent = args.agent
    env = args.env
    agent.init(
        args.net_dim,
        env.state_space.shape[0],
        env.action_space.shape[0],
        args.learning_rate,
        args.if_per_or_gae,
    )

    actor_mip = build_actor_mip(agent, env, args)
    buffer = ReplayBuffer(
        max_len=args.max_memo,
        state_dim=env.state_space.shape[0],
        action_dim=env.action_space.shape[0],
    )
    agent.state = env.reset()

    reward_record = {
        "episode": [],
        "steps": [],
        "mean_episode_reward": [],
        "unbalance": [],
        "episode_operation_cost": [],
    }
    loss_record = {
        "episode": [],
        "steps": [],
        "critic_loss": [],
        "actor_loss": [],
        "entropy_loss": [],
    }

    run = init_wandb_run(args)

    if args.train:
        while buffer.now_len < args.initial_buffer_size:
            print(f"buffer:{buffer.now_len}")
            collect_trajectory(agent, env, buffer, args, actor_mip)

        for i_episode in range(args.num_episode):
            critic_loss, actor_loss = agent.update_net(
                buffer,
                args.batch_size,
                args.repeat_times,
                args.soft_update_tau,
            )
            if actor_mip is not None:
                actor_mip.sync_network(agent.cri_target)

            run.log({"critic loss": critic_loss, "actor loss": actor_loss, "custom_step": i_episode})
            loss_record["critic_loss"].append(critic_loss)
            loss_record["actor_loss"].append(actor_loss)
            loss_record["episode"].append(i_episode)

            with torch.no_grad():
                episode_reward, episode_unbalance, episode_operation_cost = get_episode_return(
                    env,
                    agent,
                    actor_mip=actor_mip,
                )

            run.log(
                {
                    "mean_episode_reward": episode_reward,
                    "unbalance": episode_unbalance,
                    "episode_operation_cost": episode_operation_cost,
                    "custom_step": i_episode,
                }
            )
            reward_record["episode"].append(i_episode)
            reward_record["mean_episode_reward"].append(episode_reward)
            reward_record["unbalance"].append(episode_unbalance)
            reward_record["episode_operation_cost"].append(episode_operation_cost)

            print(
                f"current episode is {i_episode}, reward:{episode_reward}, "
                f"unbalance:{episode_unbalance}, buffer_length:{buffer.now_len}"
            )
            if i_episode % 10 == 0:
                agent._update_exploration_rate(args.explorate_decay, args.explorate_min)
                steps, _ = collect_trajectory(agent, env, buffer, args, actor_mip)
                reward_record["steps"].append(steps)
                loss_record["steps"].append(steps)

    run.finish()
    save_training_outputs(args.cwd, args, agent, loss_record, reward_record)
    print("training data have been saved")
    if args.save_network:
        print("training finished and actor and critic parameters have been saved")


def main():
    args = Arguments()
    for seed in args.random_seed_list:
        train_one_seed(seed, args)


if __name__ == "__main__":
    main()