dsp_bidder_3_training.py

import gymnasium as gym
import math
import random
from random import randrange
import pandas as pd
import matplotlib
import matplotlib.pyplot as plt
from collections import namedtuple, deque
from itertools import count
import numpy as np

import gymnasium as gym
from gymnasium import spaces

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

random.seed()

budget = 10
impression_max=11.888
price_max=0.118

# set up matplotlib
is_ipython = 'inline' in matplotlib.get_backend()
if is_ipython:
    from IPython import display

plt.ion()

# if GPU is to be used
device = torch.device(
    "cuda" if torch.cuda.is_available() else
    "mps" if torch.backends.mps.is_available() else
    "cpu"
)

def _normalize_vector(vector):
    if type(vector) is list:
        vector_np = np.asarray(vector, dtype=np.float32)
    else:
        vector_np = vector
    sum = np.sum(vector_np)
    if sum < 1e-8:
        return vector
    normalized_vector = vector_np / sum
    return normalized_vector


def _safe_kl(p: np.ndarray, q: np.ndarray) -> float:
    """
    KL divergence KL(p || q)
    Both p and q must be valid probability distributions.
    """
    epsilon = 0.00001
    return np.sum(p * np.log((p + epsilon) / (q + epsilon)))


def _jensen_shannon_divergence(p: np.ndarray, q: np.ndarray) -> float:
    """
    Compute Jensen–Shannon divergence between two 1D probability vectors.

    Parameters
    ----------
    p : np.ndarray
        Desired probability distribution (length 3).
    q : np.ndarray
        Current probability distribution (length 3).

    Returns
    -------
    float
        JS divergence (bounded between 0 and log(2)).
    """

    # Normalize to probability distributions
    p = _normalize_vector(p)
    q = _normalize_vector(q)

    m = 0.5 * (p + q)

    js = 0.5 * _safe_kl(p, m) + 0.5 * _safe_kl(q, m)

    return float(js)


file_screen_ids = "d:\\proj\\theneuron\\tasks\\CS_155_ml_spotzi\\005_raw_screens.csv"  # here 1500 screen ids (Strings)
df_screen_ids = pd.read_csv(file_screen_ids)
screen_ids = list(df_screen_ids['screen'])

file_inventory_last = "d:\\proj\\theneuron\\tasks\\CS_155_ml_spotzi\\013_raw_data_10dollars_publishers_venueTypes.csv" # the sample from CSV file is below:
# screen,weekday,hour,householdSmall,householdAverage,householdLarge,incomeLow,incomeAverage,incomeHigh,impressionMax,impressionHour,price,publisher1,publisher2,publisher3,venueType1,venueType2,venueType3
# 93d696ad-f4ce-4bb4-a9f1-996c771c3d7b,MONDAY,15,0.894,0.0,0.447,0.0,0.894,0.447,6.0,0.399,0.398,1.0,0.0,0.0,0.0,1.0,0.0
# 93d696ad-f4ce-4bb4-a9f1-996c771c3d7b,MONDAY,16,0.989,0.0,0.141,0.0,1.0,0.0,6.0,0.384,0.381,1.0,0.0,0.0,0.0,1.0,0.0
df_inventory = pd.read_csv(file_inventory_last)

weekdays = ['MONDAY', 'TUESDAY', 'WEDNESDAY', 'THURSDAY', 'FRIDAY', 'SATURDAY', 'SUNDAY']
hours = list(range(24))

cols = ['screen', 'weekday', 'hour']
screens_dict = {}
for (a, b, c), values in (df_inventory.set_index(cols).apply(list, axis=1)
                            .to_dict()).items():
    screens_dict.setdefault(a, {}).setdefault(b, {})[c] = values
# print(screens_dict)

def random_screen():
    return random.choice(screen_ids)


def generate_bid_requests(num_weeks):
    """Generate synthetic bid requests."""
    bid_requests = []
    for weekIndex in range(num_weeks):
        for weekday_index in range(7):
            weekday = weekdays[weekday_index]
            # print('weekday', weekday)
            for hour in hours:
                # print('  hour', hour)
                for bid_index in range(10):
                    screen_index = randrange(len(screen_ids))
                    screen_id = screen_ids[screen_index]

                    data = screens_dict[screen_id][weekday][hour]

                    householdSmall = data[0]
                    householdAverage = data[1]
                    householdLarge = data[2]
                    incomeLow = data[3]
                    incomeAverage = data[4]
                    incomeHigh = data[5]
                    impressionHour = data[7]
                    price = data[8]

                    publisher_1 = data[9]
                    publisher_2 = data[10]
                    publisher_3 = data[11]
                    venue_type_1 = data[12]
                    venue_type_2 = data[13]
                    venue_type_3 = data[14]

                    bid_request = {
                        "features": np.array([
                            # screen_index,
                            # weekday_index,
                            # hour,
                            impressionHour,
                        ], dtype=np.float32),
                        "household": np.array([
                            householdSmall,
                            householdAverage,
                            householdLarge,
                        ], dtype=np.float32),
                        "income": np.array([
                            incomeLow,
                            incomeAverage,
                            incomeHigh,
                        ], dtype=np.float32),
                        "publisher": np.array([
                            publisher_1,
                            publisher_2,
                            publisher_3,
                        ], dtype=np.float32),
                        "venue_type": np.array([
                            venue_type_1,
                            venue_type_2,
                            venue_type_3,
                        ], dtype=np.float32),
                        "price": price,
                    }
                    bid_requests.append(bid_request)
    print(f'Generated {len(bid_requests)} bid requests.')
    return bid_requests


class DspCampaign100Env(gym.Env):
    """
    Minimal DSP RL environment:
    - One episode = one campaign
    - One step = one bid request
    """

    metadata = {"render_modes": []}

    def __init__(self, bid_requests, desired_distributions, budget, impression_max, price_max):
        super().__init__()

        # ----------------------------
        # Environment data
        # ----------------------------
        self.bid_requests = bid_requests  # list of dicts (one per step)
        self.distribution_dim = 0
        for key in desired_distributions:
            dist = desired_distributions[key]
            dist2 = _normalize_vector(dist)
            desired_distributions[key] = dist2
            self.distribution_dim += len(dist2)
        self.desired_distributions = desired_distributions
        self.initial_budget = budget
        self.impression_max = impression_max
        self.price_max = price_max

        # ----------------------------
        # Action space
        # ----------------------------
        # 0 = no bid, 1 = bid
        self.action_space = spaces.Discrete(2)

        # ----------------------------
        # Observation space
        # ----------------------------
        # [current_demo(6), desired_demo(6), budget_ratio, time_ratio,
        #  bid_request_features...]
        self.bid_feat_dim = 1  # example

        obs_dim = (
            self.distribution_dim
            + 3 # campaign progress: budget_ratio, time_ratio, budget_ratio - time_ratio
            + self.bid_feat_dim
            + self.distribution_dim # bid features related to distributions (e.g. publisher, venue_type)
        )

        self.observation_space = spaces.Box(
            low=-np.inf,
            high=np.inf,
            shape=(obs_dim,),
            dtype=np.float32,
        )

        self.reset()

    # ----------------------------
    # Reset episode
    # ----------------------------
    def reset(self, seed=None, options=None):
        super().reset(seed=seed)

        self.step_idx = 0
        self.budget_left = self.initial_budget
        self.current_distributions = {}
        # self.current_demo = np.zeros(self.demo_dim, dtype=np.float32)
        for key in self.desired_distributions:
            # print("key", key, "desired_distributions[key]", type(self.desired_distributions[key]))
            self.current_distributions[key] = np.zeros(len(self.desired_distributions[key]), dtype=np.float32)

        obs = self._get_observation()
        info = {}

        return obs, info


    def reset_bid_requests(self, bid_requests):
        self.bid_requests = bid_requests


    # ----------------------------
    # Step
    # ----------------------------
    def step(self, action):
        assert self.action_space.contains(action)

        done = False

        bid = self.bid_requests[self.step_idx]

        budget_ratio = self.budget_left / self.initial_budget
        time_ratio = 1.0 - self.step_idx / len(self.bid_requests)
        cost = bid["price"] * self.price_max

        # ----------------------------
        # Apply action
        # ----------------------------
        reward = 0.0

        if action == 1 and self.budget_left >= cost:
            self.budget_left -= cost

            # Calculate distance BEFORE update
            prev_dist = 0.0
            for key in self.desired_distributions:
                prev_current = _normalize_vector(self.current_distributions[key])
                prev_dist += np.linalg.norm(self.desired_distributions[key] - prev_current)

            # Update stats
            for key in self.desired_distributions:
                self.current_distributions[key] += bid[key]

            print("desired_publishers", self.desired_distributions['publisher'], self.desired_distributions['venue_type'])
            print("current_publishers", self.current_distributions['publisher'], self.current_distributions['venue_type'])
            print("bid.publisher", bid['publisher'], "bid.venue_type", bid['venue_type'])

            # Calculate distance AFTER update
            new_dist = 0.0
            for key in self.desired_distributions:
                new_current = _normalize_vector(self.current_distributions[key])
                new_dist += np.linalg.norm(self.desired_distributions[key] - new_current)

            # --- KEY FIX ---
            # Reward is positive if we got closer to the target distribution
            # The scaling factor (e.g., * 10 or * 100) helps balance against the pacing penalty
            dist_improvement = (prev_dist - new_dist)

            # If improvement is positive (distance got smaller), good!
            # If improvement is negative (distance got bigger), bad!
            demo_reward = dist_improvement * 50.0

            reward += demo_reward
        else:
            # Pacing penalty logic remains similar, but ensure it doesn't overpower the demo reward
            # If we skip a bid, we might still be penalized if we are lagging behind budget
            pacing_gap = time_ratio - budget_ratio

            # If we are underspending (budget_ratio > time_ratio), pacing_gap is negative.
            # We punish "doing nothing" when we should be spending.
            reward += -abs(pacing_gap) * 0.5

        print("reward", reward, "action", action, "self.budget_left", self.budget_left, "time_ratio", time_ratio, "bid['price']", bid["price"] * self.price_max)

        # ----------------------------
        # Advance time
        # ----------------------------
        self.step_idx += 1

        if self.step_idx >= len(self.bid_requests) - 1:
            done = True
            # reward += - (self.budget_left / self.initial_budget) * 10

            # Optional: Final bonus for hitting the distribution target accurately
            final_dist = 0.0
            for key in self.desired_distributions:
                final_current = _normalize_vector(self.current_distributions[key])
                final_dist += np.linalg.norm(self.desired_distributions[key] - final_current)
            reward += (1.0 - final_dist) * 10  # Bonus if final distance is small

        obs = self._get_observation()
        info = {}

        return obs, reward, done, False, info

    # ----------------------------
    # Observation builder
    # ----------------------------
    def _get_observation(self):
        bid = self.bid_requests[self.step_idx]

        budget_ratio = self.budget_left / self.initial_budget
        time_ratio = 1.0 - self.step_idx / len(self.bid_requests)

        current_distributions_flat = []
        desired_distributions_flat = []
        gap_flat = []
        bid_distribution_flat = []
        for key in self.desired_distributions:
            current = self.current_distributions[key]
            current = _normalize_vector(current)
            current_distributions_flat.extend(current)

            desired = _normalize_vector(self.desired_distributions[key])
            desired_distributions_flat.extend(self.desired_distributions[key])

            gap = desired - current
            gap_flat.extend(gap.tolist())
            bid_distribution_flat.extend(bid[key])

        obs = np.concatenate([
            # np.array(current_distributions_flat, dtype=np.float32),
            # np.array(desired_distributions_flat, dtype=np.float32),
            np.array(gap_flat, dtype=np.float32),
            np.array([budget_ratio, time_ratio, budget_ratio - time_ratio], dtype=np.float32),
            bid["features"], # price
            # bid['publisher'], # publisher
            np.array(bid_distribution_flat, dtype=np.float32),
        ])
        # print("obs", obs)

        return obs.astype(np.float32)

    # ----------------------------
    # Distance metric
    # ----------------------------
    def _get_reward(self, distributions) -> float:
        """
        Compute reward based on JS divergence for household and income groups.

        Reward is NEGATIVE divergence (because we want to minimize divergence).

        Returns
        -------
        float
            Reward value (higher is better).
        """

        groups = []
        for key in self.desired_distributions:
            a = _normalize_vector(distributions[key])
            js_1 = _jensen_shannon_divergence(
                self.desired_distributions[key],
                a,
            )
            groups.append((js_1))

        total_divergence = sum(groups)

        reward = -total_divergence * 10  # scale reward to make it more significant

        return reward




# To ensure reproducibility during training, you can fix the random seeds
# by uncommenting the lines below. This makes the results consistent across
# runs, which is helpful for debugging or comparing different approaches.
#
# That said, allowing randomness can be beneficial in practice, as it lets
# the model explore different training trajectories.


seed = 42
random.seed(seed)
torch.manual_seed(seed)
if torch.cuda.is_available():
    torch.cuda.manual_seed(seed)

######################################################################
# Replay Memory
# -------------
#
# We'll be using experience replay memory for training our DQN. It stores
# the transitions that the agent observes, allowing us to reuse this data
# later. By sampling from it randomly, the transitions that build up a
# batch are decorrelated. It has been shown that this greatly stabilizes
# and improves the DQN training procedure.
#
# For this, we're going to need two classes:
#
# -  ``Transition`` - a named tuple representing a single transition in
#    our environment. It essentially maps (state, action) pairs
#    to their (next_state, reward) result, with the state being the
#    screen difference image as described later on.
# -  ``ReplayMemory`` - a cyclic buffer of bounded size that holds the
#    transitions observed recently. It also implements a ``.sample()``
#    method for selecting a random batch of transitions for training.
#

Transition = namedtuple('Transition',
                        ('state', 'action', 'next_state', 'reward'))


class ReplayMemory(object):

    def __init__(self, capacity):
        self.capacity = capacity
        self.memory = deque([], maxlen=capacity)

    def clear(self):
        self.memory = deque([], maxlen=self.capacity)

    def push(self, *args):
        """Save a transition"""
        self.memory.append(Transition(*args))

    def sample(self, batch_size):
        return random.sample(self.memory, batch_size)

    def __len__(self):
        return len(self.memory)


class DQN(nn.Module):

    def __init__(self, n_observations, n_actions):
        super(DQN, self).__init__()
        self.layer1 = nn.Linear(n_observations, 128)
        self.layer2 = nn.Linear(128, 128)
        self.layer3 = nn.Linear(128, n_actions)

    # Called with either one element to determine next action, or a batch
    # during optimization. Returns tensor([[left0exp,right0exp]...]).
    def forward(self, x):
        x = F.relu(self.layer1(x))
        x = F.relu(self.layer2(x))
        return self.layer3(x)


######################################################################
# Training
# --------
#
# Hyperparameters and utilities
# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
# This cell instantiates our model and its optimizer, and defines some
# utilities:
#
# -  ``select_action`` - will select an action according to an epsilon
#    greedy policy. Simply put, we'll sometimes use our model for choosing
#    the action, and sometimes we'll just sample one uniformly. The
#    probability of choosing a random action will start at ``EPS_START``
#    and will decay exponentially towards ``EPS_END``. ``EPS_DECAY``
#    controls the rate of the decay.
# -  ``plot_rewards`` - a helper for plotting the sum of rewards of episodes,
#    along with an average over the last 100 episodes (the measure used in
#    the official evaluations). The plot will be underneath the cell
#    containing the main training loop, and will update after every
#    episode.
#

# BATCH_SIZE is the number of transitions sampled from the replay buffer
# GAMMA is the discount factor as mentioned in the previous section
# EPS_START is the starting value of epsilon
# EPS_END is the final value of epsilon
# EPS_DECAY controls the rate of exponential decay of epsilon, higher means a slower decay
# TAU is the update rate of the target network
# LR is the learning rate of the ``AdamW`` optimizer

BATCH_SIZE = 128
# GAMMA = 0.99
# GAMMA = 0.93
GAMMA = 0.9
EPS_START = 0.9
EPS_END = 0.01
# EPS_DECAY = 2500
EPS_DECAY = 3360 / 3
# EPS_DECAY = 16800 / 3
# TAU = 0.001
# TAU = 0.005
TAU = 0.003
# LR = 1e-4
# LR = 3e-4
LR = 2e-4

# desired_household_vector = _normalize_vector([random.uniform(0, 1), random.uniform(0, 1), random.uniform(0, 1)])
# desired_income_vector = _normalize_vector([random.uniform(0, 1), random.uniform(0, 1), random.uniform(0, 1)])
# desired_publiser_vector = _normalize_vector([random.uniform(0, 1), random.uniform(0, 1), random.uniform(0, 1)])
desired_publiser_vector = _normalize_vector([0.1, 0.2, 0.7])
# desired_venue_type_vector = _normalize_vector([random.uniform(0, 1), random.uniform(0, 1), random.uniform(0, 1)])
desired_venue_type_vector = _normalize_vector([0.5, 0.3, 0.2])
env = DspCampaign100Env(generate_bid_requests(2),
                        # desired_distributions={"household": desired_household_vector, "income": desired_income_vector},
                        desired_distributions={"publisher": desired_publiser_vector, "venue_type": desired_venue_type_vector},
                        # desired_distributions={"publisher": desired_publiser_vector},
                        budget=budget, impression_max=impression_max, price_max=price_max)
# Get number of actions from gym action space
n_actions = env.action_space.n
# Get the number of state observations
state, info = env.reset()
n_observations = len(state)

policy_net = DQN(n_observations, n_actions).to(device)
target_net = DQN(n_observations, n_actions).to(device)
target_net.load_state_dict(policy_net.state_dict())

optimizer = optim.AdamW(policy_net.parameters(), lr=LR, amsgrad=True)
memory = ReplayMemory(10000)


steps_done = 0


def select_action(state):
    global steps_done
    # print('steps_done', steps_done)
    sample = random.random()
    eps_threshold = EPS_END + (EPS_START - EPS_END) * \
        math.exp(-1. * steps_done / EPS_DECAY)
    steps_done += 1
    if sample > eps_threshold:
        with torch.no_grad():
            # t.max(1) will return the largest column value of each row.
            # second column on max result is index of where max element was
            # found, so we pick action with the larger expected reward.
            return policy_net(state).max(1).indices.view(1, 1)
    else:
        return torch.tensor([[env.action_space.sample()]], device=device, dtype=torch.long)


episode_rewards = []


def plot_rewards(show_result=False):
    plt.figure(1)
    reward_t = torch.tensor(episode_rewards, dtype=torch.float)
    # print("episode_rewards", episode_rewards)
    if show_result:
        plt.title('Result')
    else:
        plt.clf()
        plt.title('Training...')
    plt.xlabel('Episode')
    plt.ylabel('Reward')
    plt.plot(reward_t.numpy())
    # Take 100 episode averages and plot them too
    # if len(reward_t) >= 100:
    #     means = reward_t.unfold(0, 100, 1).mean(1).view(-1)
    #     # print("means", means)
    #     means = torch.cat((torch.zeros(99), means))
    #     plt.plot(means.numpy())

    plt.pause(0.2)  # pause a bit so that plots are updated
    if is_ipython:
        if not show_result:
            display.display(plt.gcf())
            display.clear_output(wait=True)
        else:
            display.display(plt.gcf())


######################################################################
# Training loop
# ^^^^^^^^^^^^^
#
# Finally, the code for training our model.
#
# Here, you can find an ``optimize_model`` function that performs a
# single step of the optimization. It first samples a batch, concatenates
# all the tensors into a single one, computes :math:`Q(s_t, a_t)` and
# :math:`V(s_{t+1}) = \max_a Q(s_{t+1}, a)`, and combines them into our
# loss. By definition we set :math:`V(s) = 0` if :math:`s` is a terminal
# state. We also use a target network to compute :math:`V(s_{t+1})` for
# added stability. The target network is updated at every step with a
# `soft update <https://arxiv.org/pdf/1509.02971.pdf>`__ controlled by
# the hyperparameter ``TAU``, which was previously defined.
#

def optimize_model():
    if len(memory) < BATCH_SIZE:
        return
    transitions = memory.sample(BATCH_SIZE)
    # Transpose the batch (see https://stackoverflow.com/a/19343/3343043 for
    # detailed explanation). This converts batch-array of Transitions
    # to Transition of batch-arrays.
    batch = Transition(*zip(*transitions))

    # Compute a mask of non-final states and concatenate the batch elements
    # (a final state would've been the one after which simulation ended)
    non_final_mask = torch.tensor(tuple(map(lambda s: s is not None,
                                          batch.next_state)), device=device, dtype=torch.bool)
    non_final_next_states = torch.cat([s for s in batch.next_state
                                                if s is not None])
    state_batch = torch.cat(batch.state)
    action_batch = torch.cat(batch.action)
    reward_batch = torch.cat(batch.reward)

    # Compute Q(s_t, a) - the model computes Q(s_t), then we select the
    # columns of actions taken. These are the actions which would've been taken
    # for each batch state according to policy_net
    state_action_values = policy_net(state_batch).gather(1, action_batch)

    # Compute V(s_{t+1}) for all next states.
    # Expected values of actions for non_final_next_states are computed based
    # on the "older" target_net; selecting their best reward with max(1).values
    # This is merged based on the mask, such that we'll have either the expected
    # state value or 0 in case the state was final.
    next_state_values = torch.zeros(BATCH_SIZE, device=device)
    with torch.no_grad():
        next_state_values[non_final_mask] = target_net(non_final_next_states).max(1).values
    # Compute the expected Q values
    expected_state_action_values = (next_state_values * GAMMA) + reward_batch

    # Compute Huber loss
    criterion = nn.SmoothL1Loss()
    loss = criterion(state_action_values, expected_state_action_values.unsqueeze(1))

    # Optimize the model
    optimizer.zero_grad()
    loss.backward()
    # In-place gradient clipping
    torch.nn.utils.clip_grad_value_(policy_net.parameters(), 100)
    optimizer.step()


######################################################################
#
# Below, you can find the main training loop. At the beginning we reset
# the environment and obtain the initial ``state`` Tensor. Then, we sample
# an action, execute it, observe the next state and the reward (always
# 1), and optimize our model once. When the episode ends (our model
# fails), we restart the loop.
#
# Below, `num_episodes` is set to 600 if a GPU is available, otherwise 50
# episodes are scheduled so training does not take too long. However, 50
# episodes is insufficient for to observe good performance on CartPole.
# You should see the model constantly achieve 500 steps within 600 training
# episodes. Training RL agents can be a noisy process, so restarting training
# can produce better results if convergence is not observed.
#

if torch.cuda.is_available() or torch.backends.mps.is_available():
    num_episodes = 600
else:
    num_episodes = 250

for i_episode in range(num_episodes):
    # if i_episode == 50:
    #     memory.clear()
    # Initialize the environment and get its state
    # desired_household_vector = _normalize_vector([random.uniform(0, 1), random.uniform(0, 1), random.uniform(0, 1)])
    # desired_income_vector = _normalize_vector([random.uniform(0, 1), random.uniform(0, 1), random.uniform(0, 1)])
    # desired_publiser_vector = _normalize_vector([random.uniform(0, 1), random.uniform(0, 1), random.uniform(0, 1)])
    # desired_publiser_vector = _normalize_vector([0, 0, 1])
    # desired_venue_type_vector = _normalize_vector([random.uniform(0, 1), random.uniform(0, 1), random.uniform(0, 1)])
    # env = DspCampaign100Env(generate_bid_requests(4),
    #                         # desired_distributions={"household": desired_household_vector, "income": desired_income_vector},
    #                         # desired_distributions={"publisher": desired_publiser_vector, "venue_type": desired_venue_type_vector},
    #                         desired_distributions={"publisher": desired_publiser_vector},
    #                         budget=budget, impression_max=impression_max, price_max=price_max)
    env.reset(seed=seed)
    env.action_space.seed(seed)
    env.observation_space.seed(seed)

    sum_reward = 0
    state, info = env.reset()
    if i_episode % 3 == 0:
        env.reset_bid_requests(generate_bid_requests(2))
    state = torch.tensor(state, dtype=torch.float32, device=device).unsqueeze(0)
    for t in count():
        action = select_action(state)
        observation, reward, terminated, truncated, _ = env.step(action.item())
        if not math.isnan(reward):
            sum_reward = sum_reward + reward
        # print("sum_reward", sum_reward, "reward", reward, "terminated", terminated, "action", action)
        reward = torch.tensor([reward], device=device)
        done = terminated or truncated

        if terminated:
            next_state = None
        else:
            next_state = torch.tensor(observation, dtype=torch.float32, device=device).unsqueeze(0)

        # Store the transition in memory
        memory.push(state, action, next_state, reward)

        # Move to the next state
        state = next_state

        # Perform one step of the optimization (on the policy network)
        optimize_model()

        # Soft update of the target network's weights
        # θ′ ← τ θ + (1 −τ )θ′
        target_net_state_dict = target_net.state_dict()
        policy_net_state_dict = policy_net.state_dict()
        for key in policy_net_state_dict:
            target_net_state_dict[key] = policy_net_state_dict[key]*TAU + target_net_state_dict[key]*(1-TAU)
        target_net.load_state_dict(target_net_state_dict)

        # print("sum_reward", sum_reward)
        if done:
            # if len(episode_rewards) > 0 or sum_reward > -200:
            episode_rewards.append(sum_reward)
            plot_rewards()

            print("############# Final reward:", env._get_reward(env.current_distributions))
            print("############# Budget used:", 1 - env.budget_left / env.initial_budget)
            print("############# sum_reward:", sum_reward)
            print("############# Desire distributions:", env.desired_distributions)
            print("############# Real   distributions:", env.current_distributions)
            break

print('Complete')
plot_rewards(show_result=True)
plt.ioff()
plt.show()

MODEL_PATH = "d:\\proj\\theneuron\\tasks\\CS_155_ml_spotzi\\200_bidder_dqn_model.pt"

torch.save({
    "model_state_dict": policy_net.state_dict(),
    "optimizer_state_dict": optimizer.state_dict(),
    "n_observations": n_observations,
    "n_actions": n_actions,
}, MODEL_PATH)

print(f"Model saved to {MODEL_PATH}")