Following this instruction :
1. Strictly adhere to PEP-8 standards, ensuring code is clean, well-structured, and follows Python's style guide for readability.
2. Use appropriate modules like `typing` for type hints, ensuring code clarity, maintainability, and leveraging Python's standard libraries effectively.
3. Focus on time, space, robustness, scalability; write optimized, efficient code that can handle large datasets and scale seamlessly.
4. Implement comprehensive error handling, catching exceptions gracefully, providing meaningful error messages, and ensuring code robustness under unexpected conditions.
5. Ensure the developer agent writes advanced, maintainable Python code, balancing performance and readability while adhering to best practices.

TASK : objective we have develop below ppo  from scratch using only torch module :(assume actor , critic, any model or generalised model (may be text, video ,image , embedding wav_vec etc)

user_input :
actor :
modelconfig:
1 Model (nn.module) 
2) loss fumnction
3) optimiser

critic :
modelconfig:
1 Model (nn.module) 
2) loss fumnction
3) optimiser


Proximal Policy Optimization (PPO) is a popular algorithm for reinforcement learning that balances exploration and exploitation by updating the policy with constraints on how much it can change. PPO optimizes the following objective function by limiting the deviation from the old policy, preventing large updates that could harm learning stability.

Let's break down the main components of the PPO algorithm mathematically:

### 1. PPO Objective Function
PPO aims to maximize a clipped objective function, which is an approximation of the expected reward with constraints on policy updates. Given a policy \( \pi_\theta \) with parameters \( \theta \), PPO’s objective can be expressed as:

#### Surrogate Objective
The surrogate objective of PPO is designed to maximize the advantage while preventing large updates. Define:
- \( r_t(\theta) = \frac{\pi_\theta(a_t | s_t)}{\pi_{\theta_{\text{old}}}(a_t | s_t)} \): the probability ratio of the new policy to the old policy for action \( a_t \) at state \( s_t \).
  
- \( A_t^{\pi_{\theta_{\text{old}}}} \): the advantage function, which estimates how much better the chosen action is compared to the average action at the state.

The PPO objective function is:
\[
L^{\text{CLIP}}(\theta) = \mathbb{E}_t \left[ \min \left( r_t(\theta) A_t^{\pi_{\theta_{\text{old}}}}, \text{clip}\left(r_t(\theta), 1 - \epsilon, 1 + \epsilon\right) A_t^{\pi_{\theta_{\text{old}}}} \right) \right]
\]
where \( \epsilon \) is a hyperparameter that defines the clipping threshold (usually around 0.1 or 0.2).

This objective function consists of two terms:
1. \( r_t(\theta) A_t^{\pi_{\theta_{\text{old}}}} \): The regular surrogate objective.
2. \( \text{clip}(r_t(\theta), 1 - \epsilon, 1 + \epsilon) A_t^{\pi_{\theta_{\text{old}}}} \): The clipped surrogate objective.

The function takes the minimum of these two terms, effectively penalizing updates that result in large policy changes.

### 2. Advantage Estimation
The advantage function \( A_t^{\pi_{\theta_{\text{old}}}} \) can be calculated using methods like Generalized Advantage Estimation (GAE):
\[
A_t = \delta_t + (\gamma \lambda) \delta_{t+1} + (\gamma \lambda)^2 \delta_{t+2} + \dots
\]
where:
- \( \delta_t = r_t + \gamma V(s_{t+1}) - V(s_t) \): the Temporal Difference (TD) error.
- \( \gamma \): the discount factor.
- \( \lambda \): the GAE hyperparameter that controls bias-variance trade-off.

### 3. Policy Update
The PPO algorithm updates the policy by performing gradient ascent on the objective function \( L^{\text{CLIP}}(\theta) \):
\[
\theta \leftarrow \theta + \alpha \nabla_\theta L^{\text{CLIP}}(\theta)
\]
where \( \alpha \) is the learning rate.

### 4. Value Function Loss
In addition to the policy objective, PPO uses a value function to approximate the expected return. The loss for the value function is given by:
\[
L^{\text{VF}}(\theta) = \mathbb{E}_t \left[ \left( V_\theta(s_t) - V_t^{\text{target}} \right)^2 \right]
\]
where \( V_t^{\text{target}} \) is an estimate of the expected return at time \( t \), often computed as:
\[
V_t^{\text{target}} = \sum_{k=t}^T \gamma^{k-t} r_k
\]

### 5. Entropy Bonus
To encourage exploration, PPO includes an entropy bonus in the objective function:
\[
L^{\text{ENTROPY}}(\theta) = \mathbb{E}_t \left[ -\beta H\left( \pi_\theta(\cdot | s_t) \right) \right]
\]
where:
- \( H(\pi_\theta(\cdot | s_t)) \) is the entropy of the policy at state \( s_t \).
- \( \beta \) is a coefficient that controls the strength of the entropy bonus.

### 6. Combined Loss Function
The overall PPO loss function combines the policy loss, value function loss, and entropy bonus:
\[
L(\theta) = \mathbb{E}_t \left[ L^{\text{CLIP}}(\theta) - c_1 L^{\text{VF}}(\theta) + c_2 L^{\text{ENTROPY}}(\theta) \right]
\]
where:
- \( c_1 \) and \( c_2 \) are coefficients for the value function loss and entropy bonus, respectively.

### Summary of Key Parameters and Hyperparameters
1. \( \epsilon \): Clipping parameter for the PPO objective.
2. \( \gamma \): Discount factor for the return.
3. \( \lambda \): GAE parameter.
4. \( \alpha \): Learning rate for the policy update.
5. \( c_1 \) and \( c_2 \): Coefficients for the value loss and entropy bonus.

### PPO Algorithm Steps

1. **Collect trajectories** by running the current policy in the environment.
2. **Compute advantage estimates** \( A_t \) using GAE or another method.
3. **Compute the surrogate objective** \( L^{\text{CLIP}}(\theta) \) for policy updates.
4. **Optimize the PPO loss** \( L(\theta) \) with gradient ascent.
5. **Update the policy parameters** \( \theta \).
6. **Repeat** until convergence or a set number of iterations.

By combining these elements, PPO effectively constrains policy updates while maximizing rewards, making it a robust choice for a wide range of reinforcement learning tasks.