RL Fundamental

"""
Args:
S: list or array of states
A: list or array of actions  
P: transition probability matrix P[s][a][s'] = p(s'|s,a)
r: reward function r(s,a) or r(s,a,s')
gamma: discount factor (default 0.95)
theta: convergence threshold (default 1e-6)

Returns:
V: optimal value function
policy: optimal policy
"""

# Initialize value function
n_states = len(S)
n_actions = len(A)

# Algorithms
# ...

V = np.zeros(n_states)
while True:
    delta = 0

    # Create temporary value function for synchronous update
    V_new = np.zeros(n_states)
    for s in range(n_states):
        # Calculate Q-values for all actions
        Q = np.zeros(n_actions)
        for a in range(n_actions):
            reward = r(s, a)
            future_value = 0
            for s_next in range(n_states):
                future_value += P[s][a][s_next] * V[s_next]
            Q[a] = reward + gamma * future_value
        # Update value with maximum Q-value
        V_new[s] = np.max(Q)
        delta = max(delta, abs(V_new[s] - V[s]))
    
    # Check convergence
    V = V_new.copy()
    if delta < theta:
        break
        
# Extract optimal policy
policy = np.zeros(n_states, dtype=int)
for s in range(n_states):
    Q = np.zeros(n_actions)
    for a in range(n_actions):
        reward = r(s, a)
        future_value = 0
        for s_next in range(n_states):
            future_value += P[s][a][s_next] * V[s_next]
        Q[a] = reward + gamma * future_value
    policy[s] = np.argmax(Q)

return V, policy

policy = np.random.choice(n_actions, size=n_states)
while True:
    # Policy Evaluation
    V = np.zeros(n_states)
    while True:
        delta = 0
        V_new = np.zeros(n_states)
        for s in range(n_states):
            a = policy[s]
            reward = r(s, a)
            future_value = 0
            for s_next in range(n_states):
                future_value += P[s][a][s_next] * V[s_next]
            V_new[s] = reward + gamma * future_value
            delta = max(delta, abs(V_new[s] - V[s]))
        V = V_new.copy()
        if delta < theta:
            break

    # Policy Improvement
    policy_stable = True
    
    for s in range(n_states):
        old_action = policy[s]
        
        # Calculate Q-values for all actions
        Q = np.zeros(n_actions)
        for a in range(n_actions):
            reward = r(s, a)
            future_value = 0
            for s_next in range(n_states):
                future_value += P[s][a][s_next] * V[s_next]
            Q[a] = reward + gamma * future_value
        best_action = np.argmax(Q)
        policy[s] = best_action
        
        if old_action != best_action:
            policy_stable = False
    
    if policy_stable:
        break

return V, policy

max_eval_iters = 5
policy = np.random.choice(n_actions, size=n_states)
while True:
    # Truncated Policy Evaluation
    V = np.zeros(n_states)
    for _ in range(max_eval_iters):
        delta = 0
        V_new = np.zeros(n_states)
        for s in range(n_states):
            a = policy[s]
            reward = r(s, a)
            future_value = 0
            for s_next in range(n_states):
                future_value += P[s][a][s_next] * V[s_next]
            V_new[s] = reward + gamma * future_value
            delta = max(delta, abs(V_new[s] - V[s]))
        V = V_new.copy()
        if delta < theta:
            break

    # Policy Improvement
    policy_stable = True
    
    for s in range(n_states):
        old_action = policy[s]
        
        # Calculate Q-values for all actions
        Q = np.zeros(n_actions)
        for a in range(n_actions):
            reward = r(s, a)
            future_value = 0
            for s_next in range(n_states):
                future_value += P[s][a][s_next] * V[s_next]
            Q[a] = reward + gamma * future_value
        best_action = np.argmax(Q)
        policy[s] = best_action
        
        if old_action != best_action:
            policy_stable = False
    
    if policy_stable:
        break

return V, policy

# MC Basic
num_episodes = 100
policy = np.random.choice(n_actions, size=n_states)
while True:
    policy_stable = True
    
    # MC Policy Evaluation
    for s in range(n_states):
        old_action = policy[s]
        
        # Calculate Q-values for all actions
        Q = np.zeros(n_actions)
        for a in range(n_actions):
            Q[a] += avg_q_from_samples(s,a,num_episodes)

        # Policy Improvement
        best_action = np.argmax(Q)
        policy[s] = best_action
        
        if old_action != best_action:
            policy_stable = False
    
    if policy_stable:
        break

V = np.zeros(n_states)
for s in range(n_states):
    V[s] = np.max(Q[s])
return V, policy

# MC Exploring Starts
episodes_history
policy = np.random.choice(n_actions, size=n_states)
Q = np.zeros((n_states, n_actions))
returns = {}

# Policy Evaluation
for episode_history in episodes_history:
    # Calculate Q-values for all actions
    G = 0 
    visited = set()
    for t in range(len(episode_history)-1, -1, -1):
        state, action, reward = episode_history[t]
        state_action = (state, action)
        G = reward + gamma * G
        
        if state_action not in visited:
            visited.add(state_action)
            
            if state_action not in returns:
                returns[state_action] = []
            returns[state_action].append(G)
            
            Q[state, action] = np.mean(returns[state_action])
    
    # Policy Improvement
    for s in range(n_states):
        best_action = np.argmax(Q[s])
        policy[s] = best_action
        
V = np.zeros(n_states)
for s in range(n_states):
    V[s] = np.max(Q[s])
return V, policy

# MC xi-Greedy
xi = 0.1
episodes_history
policy = np.random.choice(n_actions, size=n_states)
Q = np.zeros((n_states, n_actions))
returns = {}

# Policy Evaluation
for episode_history in episodes_history:
    # Calculate Q-values for all actions
    G = 0 
    for t in range(len(episode_history)-1, -1, -1):
        state, action, reward = episode_history[t]
        state_action = (state, action)
        G = reward + gamma * G
            
        if state_action not in returns:
            returns[state_action] = []
		returns[state_action].append(G)
            
        Q[state, action] = np.mean(returns[state_action])
    
    # ε-greedy Policy Improvement
    for s in range(n_states):
        if np.random.random() > epsilon:
            best_action = np.argmax(Q[s])
            policy[s] = best_action
        else:
            policy[s] = np.random.choice(n_actions)
        
V = np.zeros(n_states)
for s in range(n_states):
    V[s] = np.max(Q[s])
return V, policy

def get_best_action(state):
    # ε-greedy Policy Improvement
	if np.random.random() > epsilon:
        best_action = np.argmax(Q[state])
    else:
        best_action = np.random.choice(n_actions)
    return best_action

trajectory = []
total_reward = 0
Q = np.zeros((n_states, n_actions))

cur_state = start_state
cur_action = get_best_action(cur_state)
trajectory.append(0, cur_state, cur_action)

while cur_state is not target_state:
	reward = r(cur_state, cur_action)
    next_state = np.random.choice(n_states, p=P[cur_state][cur_action])
    next_action = get_best_action(next_state)
    total_reward += reward
    trajectory.append(reward,next_state, next_action)
    
    # TD Policy Evaluation  
    Q[cur_state][cur_action] = Q[cur_state][cur_action] + alpha * (reward + gamma * Q[next_state][next_action] -  Q[cur_state][cur_action])
    cur_state, cur_action = next_state, next_action

policy = np.argmax(Q, axis=1)
V = np.zeros(n_states)
for s in range(n_states):
    V[s] = np.max(Q[s])
return V, policy

def get_best_action(state):
    # ε-greedy Policy Improvement
	if np.random.random() > epsilon:
        action = np.argmax(Q[state])
    else:
        action = np.random.choice(n_actions)
    return action
def get_expected_q_value(state):
    # Expected Q
    best_action = np.argmax(Q[state])
    action_probs = np.ones(n_actions) * (epsilon / n_actions) 
    action_probs[best_action] += (1 - epsilon)
    
    expected_value = 0
    for a in range(n_actions):
        expected_value += action_probs[a] * Q[state][a]
    return expected_value

trajectory = []
total_reward = 0
Q = np.zeros((n_states, n_actions))

cur_state = start_state
cur_action = get_best_action(cur_state)
trajectory.append(0, cur_state, cur_action)

while cur_state is not target_state:
	reward = r(cur_state, cur_action)
    next_state = np.random.choice(n_states, p=P[cur_state][cur_action])
    next_action = get_best_action(next_state)
    total_reward += reward
    trajectory.append(reward,next_state, next_action)
    
    # TD Policy Evaluation  
    Q[cur_state][cur_action] = Q[cur_state][cur_action] + alpha * (reward + gamma * get_expected_q_value(next_state) -  Q[cur_state][cur_action])
    cur_state, cur_action = next_state, next_action

policy = np.argmax(Q, axis=1)
V = np.zeros(n_states)
for s in range(n_states):
    V[s] = np.max(Q[s])
return V, policy

def get_best_action(state):
    # ε-greedy Policy Improvement
	if np.random.random() > epsilon:
        best_action = np.argmax(Q[state])
    else:
        best_action = np.random.choice(n_actions)
    return best_action

trajectory = []
total_reward = 0
Q = np.zeros((n_states, n_actions))

cur_state = start_state
cur_action = get_best_action(cur_state)
trajectory.append(0, cur_state, cur_action)

n_steps
experience_buffer = []

while cur_state is not target_state:
	reward = r(cur_state, cur_action)
    next_state = np.random.choice(n_states, p=P[cur_state][cur_action])
    next_action = get_best_action(next_state)
    total_reward += reward
    trajectory.append(reward,next_state, next_action)
    
    # n-steps
    experience = (cur_state, cur_action, reward)
    experience_buffer.append(experience)
    
    # TD Policy Evaluation
    if len(experience_buffer) >= n_steps:
        n_step_ago_state, n_step_ago_action, n_step_reward = experience_buffer[0]
        
        n_step_return = 0
        for i in range(len(experience_buffer)):
            n_step_return += (gamma ** i) * experience_buffer[i][2]
        if next_state is not target_state:
            n_step_return += (gamma ** n_steps) * Q[next_state][next_action]
            
        Q[n_step_ago_state][n_step_ago_action] = Q[n_step_ago_state][n_step_ago_action] + alpha * (n_step_return -  Q[n_step_ago_state][n_step_ago_action])
        experience_buffer.pop(0)
        
    cur_state, cur_action = next_state, next_action

# left n-steps
while len(experience_buffer) > 0:
        n_step_ago_state, n_step_ago_action, n_step_reward = experience_buffer[0]
        
        n_step_return = 0
        for i in range(len(experience_buffer)):
            n_step_return += (gamma ** i) * experience_buffer[i][2]
        
        Q[n_step_ago_state][n_step_ago_action] = Q[n_step_ago_state][n_step_ago_action] + alpha * (
            n_step_return - Q[n_step_ago_state][n_step_ago_action]
        )
        
        experience_buffer.pop(0)
    
policy = np.argmax(Q, axis=1)
V = np.zeros(n_states)
for s in range(n_states):
    V[s] = np.max(Q[s])
return V, policy

# Q-learning off-policy
episodes_history
Q = np.zeros((n_states, n_actions))

# Policy Evaluation
for episode_history in episodes_history:
    for t in range(len(episode_history)-1):
        state, action, reward = episode_history[t]
        next_state, next_action, next_reward = episode_history[t+1]
        
        if t == len(episode_history)-2: 
            q_target = reward + gamma * next_reward
        else:
            q_target = reward + gamma * np.max(Q[next_state])
        Q[state, action] = Q[state, action] + alpha * (q_target - Q[state, action])

# Policy Improvement
policy = np.random.choice(n_actions, size=n_states)
for s in range(n_states):
    best_action = np.argmax(Q[s])
    policy[s] = best_action

V = np.zeros(n_states)
for s in range(n_states):
    V[s] = np.max(Q[s])
return V, policy

# DQN off-policy
class DQN(nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, 128)
        self.fc2 = nn.Linear(128, 128)
        self.fc3 = nn.Linear(128, output_dim)
        
    def forward(self, x):
        x1 = torch.relu(self.fc1(x))
        x2 = torch.relu(self.fc2(x1))
        return self.fc3(x2)

class DQNAgent:
    def __init__(self, state_dim, action_dim, learning_rate, gamma):
        self.state_dim = state_dim
        self.action_dim = action_dim
        
        # Initialize main network and target network
        self.q_network = DQN(state_dim, action_dim)
        self.target_network = DQN(state_dim, action_dim)
        self.target_network.load_state_dict(self.q_network.state_dict())
        self.optimizer = optim.Adam(self.q_network.parameters(), lr=learning_rate)
        
        self.batch_size = 32
        self.memory_size = 10000
        self.memory = deque(maxlen=self.memory_size)
        self.learning_step = 0
        self.target_update_freq = 100

    def get_best_action(self, state):
        # ε-greedy Policy Improvement for nn
        if np.random.random() > epsilon:
            with torch.no_grad():
                state_tensor = torch.FloatTensor(state).unsqueeze(0)
                q_values = self.q_network(state_tensor)
                best_action = np.argmax(q_values.numpy())
        else:
            best_action = np.random.choice(self.action_dim)
        return best_action
    
    def store_experience(self, state, action, reward, next_state, done):
        self.memory.append((state, action, reward, next_state, done))
    
    def update_network(self):
        if len(self.memory) < self.batch_size:
            return
        
        # Sample batch from experience replay
        batch = random.sample(self.memory, self.batch_size)
        states, actions, rewards, next_states, dones = zip(*batch)
        states = torch.FloatTensor(states)
        actions = torch.LongTensor(actions).unsqueeze(1)
        rewards = torch.FloatTensor(rewards)
        next_states = torch.FloatTensor(next_states)
        dones = torch.BoolTensor(dones)
        
        # Current Q values
        current_q_values = self.q_network(states).gather(1, actions).squeeze()
        
        # TD target
        with torch.no_grad():
            next_q_values = self.target_network(next_states).max(1)[0]
            # Set Q value to 0 for terminal states
            next_q_values[dones] = 0.0
            target_q_values = rewards + gamma * next_q_values
        
        # Gradient Descent
        loss = nn.MSELoss()(current_q_values, target_q_values)
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
        
        # Update target network periodically
        self.learning_step += 1
        if self.learning_step % self.target_update_freq == 0:
            self.target_network.load_state_dict(self.q_network.state_dict())
            
        return loss.item()

def train_dqn(env, n_episodes=1000):
    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.n
    agent = DQNAgent(state_dim, action_dim,learning_rate, gamma)
    
    for episode in range(n_episodes):
        total_reward = 0
        
        # Environment reset
        cur_state = env.reset()
        cur_action = agent.get_best_action(cur_state)
        done = False
        step_count = 0
        while not done:
            # Take action and observe next state and reward
            next_state, reward, done, truncated, _ = env.step(cur_action)
            next_action = agent.get_best_action(next_state)
            total_reward += reward
            done = done or truncated
            
            agent.store_experience(cur_state, cur_action, reward, next_state, done)
            loss = agent.update_network()
            
            cur_state, cur_action = next_state, next_action
            step_count += 1
        
        if episode % 100 == 0:
            print(f"Episode {episode}, Total Reward: {total_reward:.2f}, Steps: {step_count}, Loss: {loss}")
            
	return agent

# REINFORCE on policy
class PolicyNetwork(nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, 128)
        self.fc2 = nn.Linear(128, 128)
        self.fc3 = nn.Linear(128, output_dim)
        
    def forward(self, x):
        x1 = torch.relu(self.fc1(x))
        x2 = torch.relu(self.fc2(x1))
        # Use softmax to output action probabilities
        return torch.softmax(self.fc3(x2), dim=-1)

class REINFORCEAgent:
    def __init__(self, state_dim, action_dim, learning_rate, gamma):
        self.state_dim = state_dim
        self.action_dim = action_dim
        
        # Initialize policy network
        self.policy_network = PolicyNetwork(state_dim, action_dim)
        self.optimizer = optim.Adam(self.policy_network.parameters(), lr=learning_rate)
        
        self.trajectory = []
    
    def get_best_action(self, state, training=True):
        state_tensor = torch.FloatTensor(state).unsqueeze(0)
        action_probs = self.policy_network(state_tensor)
        
        if training:
            # Sample action from probability distribution during training
            action_dist = torch.distributions.Categorical(action_probs)
            action = action_dist.sample()
            log_prob = action_dist.log_prob(action)
            return action.item(), log_prob
        else:
            # Use greedy action during evaluation
            return torch.argmax(action_probs).item(), None
    
      def get_policy(self, states):
        """Extract policy"""
        policy = []
        
        with torch.no_grad():
            for state in states:
                state_tensor = torch.FloatTensor(state).unsqueeze(0)
                action_probs = self.policy_network(state_tensor).numpy()[0]
                best_action = np.argmax(action_probs)
                policy.append(best_action)
        
        return np.array(policy)
    
    def store_transition(self, state, action, log_prob, reward, done):
        self.trajectory.append((state, action, log_prob, reward, done))
    
    def update_policy(self):
        if not self.trajectory:
            return 0
        *_, log_probs, rewards, _ = zip(*self.trajectory)
        
        # Caculate returns
        returns = []
        G = 0
        for r in reversed(rewards):
            G = r + gamma * G
            returns.insert(0, G)
        returns = torch.FloatTensor(returns)
        returns = (returns - returns.mean()) / (returns.std() + 1e-8)
        
        log_probs = torch.stack(log_probs)
        
        # Gradient Rise
        loss = (-log_probs * returns).mean()
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
        
        self.trajectory = []
        
        return loss.item()

def train_reinforce(env, n_episodes=1000):
    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.n
    agent = REINFORCEAgent(state_dim, action_dim, learning_rate, gamma)
    
    for episode in range(n_episodes):
        total_reward = 0
        
        # Environment reset
        state = env.reset()
        done = False
        step_count = 0
        while not done:
            action, log_prob = agent.get_best_action(state)
            next_state, reward, done, truncated, _  = env.step(action)
            total_reward += reward
            done = done or truncated
            
            agent.store_transition(state, action, log_prob, reward, done)
            loss = agent.update_policy()
            
            state = next_state
            step_count += 1
                
        if episode % 100 == 0:
            print(f"Episode {episode}, Total Reward: {total_reward:.2f}, Steps: {step_count}, Loss: {loss}")

	return agent

# A2C on policy
class ActorCriticNetwork(nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        # Shared feature extraction layers
        self.shared_fc1 = nn.Linear(input_dim, 128)
        self.shared_fc2 = nn.Linear(128, 128)
        
        # Actor head - outputs action probabilities
        self.actor_fc = nn.Linear(128, output_dim)
        
        # Critic head - outputs state value
        self.critic_fc = nn.Linear(128, 1)
        
    def forward(self, x):
        x1 = torch.relu(self.shared_fc1(x))
        x2 = torch.relu(self.shared_fc2(x1))
        
        # Actor output: action probabilities
        action_probs = torch.softmax(self.actor_fc(x2), dim=-1)
        
        # Critic output: state value
        state_value = self.critic_fc(x2)
        
        return action_probs, state_value

class A2CAgent:
    def __init__(self, state_dim, action_dim, learning_rate, gamma):
        self.state_dim = state_dim
        self.action_dim = action_dim
        
        self.entropy_coef = 0.01
        
        # Initialize actor-critic network
        self.network = ActorCriticNetwork(state_dim, action_dim)
        self.optimizer = optim.Adam(self.network.parameters(), lr=learning_rate)
        
        self.trajectory = []
    
    def get_best_action(self, state, training=True):
        state_tensor = torch.FloatTensor(state).unsqueeze(0)
        action_probs, state_value = self.network(state_tensor)
        
        if training:
            # Sample action from probability distribution during training
            action_dist = torch.distributions.Categorical(action_probs)
            action = action_dist.sample()
            log_prob = action_dist.log_prob(action)
            entropy = action_dist.entropy()
            return action.item(), log_prob, entropy, state_value.squeeze()
        else:
            # Use greedy action during evaluation
            return torch.argmax(action_probs).item(), None, None, state_value.squeeze()
    
    def get_policy(self, states):
        """Extract policy"""
        policy = []
        
        with torch.no_grad():
            for state in states:
                state_tensor = torch.FloatTensor(state).unsqueeze(0)
                action_probs, _ = self.network(state_tensor)
                action_probs = action_probs.numpy()[0]
                best_action = np.argmax(action_probs)
                policy.append(best_action)
        
        return np.array(policy)
    
    def store_transition(self, state, action, log_prob, entropy, value, reward, done):
        self.trajectory.append((state, action, log_prob, entropy, value, reward, done))
    
    def update_policy(self):
        if not self.trajectory:
            return 0, 0, 0
        states, *_, log_probs, entropies, values, rewards, dones = zip(*self.trajectory)
        
        # Calculate advantages and returns
        advantages = []
        returns = []
    	G = 0
        for i in reversed(range(len(rewards))):
            G = rewards[i] + gamma * G * (1 - dones[i])
            returns.insert(0, G)
            advantages.insert(0, G - values[i]) 
        advantages = torch.FloatTensor(advantages)
        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
        returns = torch.FloatTensor(returns)
        
		log_probs = torch.stack(log_probs)
        values_tensor = torch.stack(values)
        entropies = torch.stack(entropies)

        # Calculate Rise
        actor_loss = (-log_probs * advantages.detach()).mean()
        # actor_loss_off_policy = (-(self.network(states)/log_probs)*(log_probs * advantages.detach())).mean()
        critic_loss = 0.5 * (returns - values_tensor).pow(2).mean()
        entropy_loss = -entropies.mean()
        total_loss = actor_loss + critic_loss + self.entropy_coef * entropy_loss
        self.optimizer.zero_grad()
        total_loss.backward()
        torch.nn.utils.clip_grad_norm_(self.parameters(), max_norm=0.5)
        self.optimizer.step()

        self.trajectory = []

        return actor_loss.item(), critic_loss.item(), entropy_loss.item()

def train_a2c(env, n_episodes=1000):
    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.n
    agent = A2CAgent(state_dim, action_dim, learning_rate, gamma)
    
    for episode in range(n_episodes):
        total_reward = 0
        
        # Environment reset
        state = env.reset()
        done = False
        step_count = 0
        while not done:
            action, log_prob, entropy, value = agent.get_best_action(state)
            next_state, reward, done, truncated, _ = env.step(action)
            total_reward += reward
            done = done or truncated
            
            agent.store_transition(state, action, log_prob, entropy, value, reward, done)
            actor_loss, critic_loss, entropy_loss = agent.update_policy()
            
            state = next_state
            step_count += 1
        
        if episode % 100 == 0:
            print(f"Episode {episode}, Total Reward: {total_reward:.2f}, Steps: {step_count}")
            print(f"  Actor Loss: {actor_loss:.4f}, Critic Loss: {critic_loss:.4f}, Entropy Loss: {entropy_loss:.4f}")

def test_a2c(env, agent):
    state = env.reset()
    done = False
    total_reward = 0
    
    while not done:
        action, _, _, _ = agent.get_best_action(state, training=False)
        next_state, reward, done, truncated, _ = env.step(action)
        done = done or truncated
        total_reward += reward
        state = next_state
    
    print(f"Test Reward: {total_reward}")

# DPG off policy
# policy update net
class ActorNetwork(nn.Module):
    def __init__(self, input_dim, output_dim, action_high, action_low):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, 128)
        self.fc2 = nn.Linear(128, 128)
        self.actor_fc = nn.Linear(128, output_dim)
        
        # Action scaling parameters.
        self.action_high = torch.FloatTensor(action_high)
        self.action_low = torch.FloatTensor(action_low)
        self.action_range = self.action_high - self.action_low
        
    def forward(self, x):
        x1 = torch.relu(self.fc1(x))
        x2 = torch.relu(self.fc2(x1))
        
        # Scale to environment's action range
        action = torch.tanh(self.actor_fc(x2))
        scaled_action = self.action_low + (action + 1.0) * 0.5 * self.action_range
        return scaled_action

# policy evaluation net
class CriticNetwork(nn.Module):
    def __init__(self, state_dim, action_dim):
        super().__init__()
        self.fc1 = nn.Linear(state_dim + action_dim, 128)
        self.fc2 = nn.Linear(128, 128)
        self.critic_fc = nn.Linear(128, 1)
        
    def forward(self, state, action):
        x1 = torch.cat([state, action], dim=-1)
        x2 = torch.relu(self.fc1(x1))
        x3 = torch.relu(self.fc2(x2))
        q_value = self.critic_fc(x3)
        return q_value

class DPGAgent:
    def __init__(self, state_dim, action_dim, action_high, action_low, 
                 learning_rate_actor, learning_rate_critic, 
                 gamma, tau):
        self.state_dim = state_dim
        self.action_dim = action_dim
        self.tau = tau  # Soft update parameter
        
        # Initialize networks
        self.actor = ActorNetwork(state_dim, action_dim, action_high, action_low)
        self.critic = CriticNetwork(state_dim, action_dim)

        # Initialize target networks
        self.actor_target = ActorNetwork(state_dim, action_dim, action_high, action_low)
        self.critic_target = CriticNetwork(state_dim, action_dim)

        # Copy weights to target networks
        self.hard_update(self.actor_target, self.actor)
        self.hard_update(self.critic_target, self.critic)
        
        # Optimizers
        self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=learning_rate_actor)
        self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=learning_rate_critic)

        # Experience replay buffer
        self.batch_size = 64
        self.memory_size = 10000
        self.memory = deque(maxlen=self.memory_size)
         self.learning_step = 0
        self.target_update_freq = 100
        
        # For exploration
        self.noise_std = 0.1
        self.noise_decay = 0.999
        
    def hard_update(self, target, source):
        """Copy weights from source to target network"""
        for target_param, param in zip(target.parameters(), source.parameters()):
            target_param.data.copy_(param.data)

    def soft_update(self, target, source):
        """Soft update target network parameters"""
        for target_param, param in zip(target.parameters(), source.parameters()):
            target_param.data.copy_(self.tau * param.data + (1.0 - self.tau) * target_param.data)
    
    def get_action(self, state, training=True):
        state_tensor = torch.FloatTensor(state).unsqueeze(0)
        action = self.actor(state_tensor).detach().numpy()[0]
        
        if training:
            # Add exploration noise
            noise = np.random.normal(0, self.noise_std, size=self.action_dim)
            action = np.clip(action + noise, self.actor.action_low.numpy(), self.actor.action_high.numpy())
        
        return action
    
    def store_experience(self, state, action, reward, next_state, done):
        self.memory.append((state, action, reward, next_state, done))
    
    def update_policy(self):
        if len(self.memory) < self.batch_size:
            return 0, 0
        
        # Sample batch from replay buffer
        batch = random.sample(self.memory, self.batch_size)
        states, actions, rewards, next_states, dones = zip(*batch)
        states = torch.FloatTensor(states)
        actions = torch.FloatTensor(actions)
        rewards = torch.FloatTensor(rewards).unsqueeze(1)
        next_states = torch.FloatTensor(next_states)
        dones = torch.BoolTensor(dones).unsqueeze(1)
        
        # Update Critic
        with torch.no_grad():
            next_actions = self.actor_target(next_states)
            target_q = self.critic_target(next_states, next_actions)
            target = rewards + gamma * target_q * (~dones).float()
        
        current_q = self.critic(states, actions)
        critic_loss = nn.MSELoss()(current_q, target)
        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        torch.nn.utils.clip_grad_norm_(self.critic.parameters(), max_norm=0.5)
        self.critic_optimizer.step()
        
        # Update Actor
        policy_actions = self.actor(states)
        actor_loss = -self.critic(states, policy_actions).mean()
        self.actor_optimizer.zero_grad()
        actor_loss.backward()
        torch.nn.utils.clip_grad_norm_(self.actor.parameters(), max_norm=0.5)
        self.actor_optimizer.step()
        
        # Soft update target networks
        self.soft_update(self.actor_target, self.actor)
        self.soft_update(self.critic_target, self.critic)
        
        self.learning_step += 1
        # Hard update target networks
        if self.learning_step % self.target_update_freq == 0:
            self.hard_update(self.actor_target, self.actor)
            self.hard_update(self.critic_target, self.critic)
        
        # Decay exploration noise
        self.noise_std *= self.noise_decay
        
        return actor_loss.item(), critic_loss.item()
    
    def get_policy(self, states):
        """Extract deterministic policy"""
        policy = []
        
        with torch.no_grad():
            for state in states:
                state_tensor = torch.FloatTensor(state).unsqueeze(0)
                action = self.actor(state_tensor).numpy()[0]
                policy.append(action)
        
        return np.array(policy)

def train_dpg(env, n_episodes=1000):
    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.shape[0]
    action_high = env.action_space.high
    action_low = env.action_space.low
    agent = DPGAgent(state_dim, action_dim, action_high, action_low)
    
    for episode in range(n_episodes):
        total_reward = 0
        
        state = env.reset()
        done = False
        step_count = 0
        while not done:
            action = agent.get_action(state)
            next_state, reward, done, truncated, _ = env.step(action)
            total_reward += reward
            done = done or truncated
            
            agent.store_experience(state, action, reward, next_state, done)
            actor_loss, critic_loss = agent.update_policy()
            
            state = next_state
            step_count += 1
        
        if episode % 100 == 0:
            print(f"Episode {episode}, Total Reward: {total_reward:.2f}, Steps: {step_count}")
            print(f"  Actor Loss: {actor_loss:.4f}, Critic Loss: {critic_loss:.4f}")
            print(f"  Noise STD: {agent.noise_std:.4f}")
            
	return agent

def test_dpg(env, agent):
    state = env.reset()
    done = False
    total_reward = 0
    
    while not done:
        action = agent.get_action(state, training=False)  # No exploration during test
        next_state, reward, done, truncated, _ = env.step(action)
        done = done or truncated
        total_reward += reward
        state = next_state
    
    print(f"Test Reward: {total_reward}")