Final Project： Model-Based and Reinforcement Learning Control of the Acrobot

# Install dependencies for Colab
%pip install -q control gymnasium stable-baselines3[extra] scipy matplotlib

import numpy as np
import matplotlib.pyplot as plt

# For numerical integration
try:
    import scipy as sp
    from scipy.integrate import solve_ivp
except ImportError:
    %pip install scipy
    import scipy as sp
    from scipy.integrate import solve_ivp

# For control systems
try:
    import control as ctl
    print(f"control already installed: {ctl.__version__}")
except ImportError:
    %pip install control
    import control as ctl

# For RL APIs
try:
    import gymnasium as gym
    print(f"gymnasium already installed: {gym.__version__}")
except ImportError:
    %pip install gymnasium
    import gymnasium as gym

# For RL algorithms
try:
    import stable_baselines3 as sb3
    print(f"stable_baselines3 already installed: {sb3.__version__}")
except ImportError:
    %pip install stable-baselines3
    import stable_baselines3 as sb3

# For visualization
try:
    import pygame
    print(f"pygame already installed: {pygame.__version__}")
except ImportError:
    %pip install pygame

# Set default plotting style
plt.style.use("ggplot")

# Example default parameters (you can adjust)
default_params = dict(
    m1=1.0, m2=1.0,
    l1=1.0, l2=1.0,
    lc1=0.5, lc2=0.5,
    I1=1.0, I2=1.0,
    g=9.81
)

def acrobot_dynamics(t, x, u, params):
    """
    Continuous-time Acrobot dynamics:
        x = [theta1, theta2, theta1_dot, theta2_dot]
        u = torque at joint 2 (scalar)
    Returns xdot.
    """
    # Unpack state
    th1, th2, omega1, omega2 = x

    # Unpack parameters
    m1, m2, l1, lc1, lc2, I1, I2, g = (params.get(key) for key in ['m1', 'm2', 'l1', 'lc1', 'lc2', 'I1', 'I2', 'g'])

    # Compute intermediate terms
    d1 = (m1 * lc1**2 +
          m2 * (l1**2 + lc2**2 + 2 * l1 * lc2 * np.cos(th2)) +
          I1 + I2)

    d2 = m2 * (lc2**2 + l1 * lc2 * np.cos(th2)) + I2

    phi2 = m2 * lc2 * g * np.cos(th1 + th2 - np.pi/2)

    phi1 = (-m2 * l1 * lc2 * omega2**2 * np.sin(th2)
            - 2 * m2 * l1 * lc2 * omega1 * omega2 * np.sin(th2)
            + (m1 * lc1 + m2 * l1) * g * np.cos(th1 - np.pi/2)
            + phi2)

    # Compute angular accelerations
    # From the coupled equations, solve for th2_dd first, then th1_dd
    denominator = m2 * lc2**2 + I2 - d2**2 / d1
    th2_dd = (u + (d2 / d1) * phi1
              - m2 * l1 * lc2 * omega1**2 * np.sin(th2)
              - phi2) / denominator

    th1_dd = -(d2 * th2_dd + phi1) / d1

    # Return state derivative
    return np.array([omega1, omega2, th1_dd, th2_dd])

def simulate_acrobot(x0, u_fun, T=5.0, params=default_params):
    """
    Simulate Acrobot with given input function u_fun(t, x).
    """
    def fwrap(t, x):
        u = u_fun(t, x)
        return acrobot_dynamics(t, x, u, params)

    sol = solve_ivp(fwrap, [0, T], x0, max_step=0.01, rtol=1e-6, atol=1e-8)
    return sol.t, sol.y

# Example usage (uncomment after implementing dynamics):
x0 = np.array([0.3, 0.4, 0.0, 0.0])
t, X = simulate_acrobot(x0, lambda t, x: 0.0) # lambda t, x: 0.0 means u=0
plt.figure()
plt.plot(t, X[0], label="theta1")
plt.plot(t, X[2], label="theta2")
plt.legend()
plt.xlabel("Time [s]")
plt.ylabel("Angle [rad]")
plt.title("Passive Acrobot Dynamics (u=0)")
plt.show()

from matplotlib.animation import FuncAnimation
from IPython.display import HTML, display

def animate_acrobot_solution(sol, x0, params=None, duration=5.0):
    """
    Create animation from a solved acrobot ODE solution.

    Parameters:
    sol: Solution object containing time array t and state array y
    x0: Initial state [theta1, theta2, omega1, omega2]
    params: Dictionary of acrobot parameters (uses default_params if None)
    duration: Animation duration (s)

    Returns:
    HTML animation object
    """
    if params is None:
        params = default_params

    l1 = params['l1']
    l2 = params['l2']

    # Set up the figure and axis for animation
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 6))

    # Acrobot animation subplot
    total_length = l1 + l2
    ax1.set_xlim(-total_length * 1.2, total_length * 1.2)
    ax1.set_ylim(-total_length * 1.2, total_length * 1.2)
    ax1.set_aspect('equal')
    ax1.set_title(f'Acrobot Animation (θ₁₀={x0[0]:.2f}, θ₂₀={x0[1]:.2f})')
    ax1.grid(True)
    ax1.plot(0, 0, 'ko', markersize=8)  # Fixed pivot point

    # Create acrobot elements
    link1, = ax1.plot([], [], 'b-', linewidth=4, label='Link 1')
    link2, = ax1.plot([], [], 'r-', linewidth=4, label='Link 2')
    joint1, = ax1.plot([], [], 'ko', markersize=8)
    joint2, = ax1.plot([], [], 'go', markersize=10)
    trace, = ax1.plot([], [], 'g--', alpha=0.3, linewidth=1)
    ax1.legend()

    # Time series subplot
    ax2.set_xlim(0, duration)
    ax2.set_ylim(-np.pi, np.pi)
    ax2.set_xlabel('Time (s)')
    ax2.set_ylabel('Angle (rad)')
    ax2.set_title('Joint Angles vs Time')
    ax2.grid(True)

    theta1_line, = ax2.plot([], [], 'b-', label='θ₁(t)', linewidth=2)
    theta2_line, = ax2.plot([], [], 'r-', label='θ₂(t)', linewidth=2)
    time_marker1, = ax2.plot([], [], 'bo', markersize=6)
    time_marker2, = ax2.plot([], [], 'ro', markersize=6)
    ax2.legend()

    # Store trace points
    x_trace, y_trace = [], []

    # Animation function
    def animate(frame):
        if frame < len(sol.t):
            # Current state
            theta1 = sol.y[0][frame]
            theta2 = sol.y[1][frame]
            t_current = sol.t[frame]

            # Joint positions
            x1 = l1 * np.sin(theta1)
            y1 = -l1 * np.cos(theta1)
            x2 = x1 + l2 * np.sin(theta1 + theta2)
            y2 = y1 - l2 * np.cos(theta1 + theta2)

            # Update acrobot links
            link1.set_data([0, x1], [0, y1])
            link2.set_data([x1, x2], [y1, y2])
            joint1.set_data([x1], [y1])
            joint2.set_data([x2], [y2])

            # Update trace (last 100 points)
            x_trace.append(x2)
            y_trace.append(y2)
            if len(x_trace) > 100:
                x_trace.pop(0)
                y_trace.pop(0)
            trace.set_data(x_trace, y_trace)

            # Update time series
            theta1_line.set_data(sol.t[:frame+1], sol.y[0][:frame+1])
            theta2_line.set_data(sol.t[:frame+1], sol.y[1][:frame+1])
            time_marker1.set_data([t_current], [theta1])
            time_marker2.set_data([t_current], [theta2])

        return link1, link2, joint1, joint2, trace, theta1_line, theta2_line, time_marker1, time_marker2

    # Create and display animation
    anim = FuncAnimation(fig, animate, frames=len(sol.t), interval=50, blit=True, repeat=True)
    plt.tight_layout()
    plt.show()

    # Return HTML version for display
    return HTML(anim.to_jshtml())

# Create solution object from existing variables
from types import SimpleNamespace
solution = SimpleNamespace(t=t, y=X)

print(f"Animating acrobot with initial state: θ₁={x0[0]:.2f}, θ₂={x0[1]:.2f}, ω₁={x0[2]:.2f}, ω₂={x0[3]:.2f}")
print(f"Solution contains {len(solution.t)} time points over {solution.t[-1]:.2f} seconds")

# Create and display animation
print("Creating animation...")
animation_html = animate_acrobot_solution(solution, x0, default_params, duration=t[-1])
display(animation_html)

# Use the default parameters defined earlier for the Acrobot model
acrobot_params = default_params

# State derivative function for python-control
def acrobot_update(t, x, u, params):
    """
    Wrapper for the acrobot_dynamics function to be used with control.nlsys.
    The input 'u' from nlsys is an array, so we extract the scalar torque.
    """
    # TODO: Call the acrobot_dynamics function with the correct inputs.

# Output function for python-control
def acrobot_output(t, x, u, params):
    """
    Defines the system output. For state-feedback control, we assume we have the full state as the output.
    """
    # TODO: Return the state vector as the output.


# Create the nonlinear system object using nlsys
acrobot_nlsys = ctl.nlsys(
    acrobot_update,
    acrobot_output,
    name='acrobot',
    params=acrobot_params,
    states=['theta1', 'theta2', 'omega1', 'omega2'],
    outputs=['theta1', 'theta2', 'omega1', 'omega2'],
    inputs=['tau']
)

# After filling in the TODOs, linearize the system using acrobot_nlsys.linearize
# TODO: linearize the system at the upright equilibrium point

# TODO: choose equilibrium and design LQR

# TODO: write your code here for the optimal trajecotry generation

def acrobot_err_update(t, e, du, params):
    # e: error state (4,)
    # du: error input (1,)  (delta torque)
    xref = params['x_ref'](t).ravel()    # x_ref(t)
    uref = params['u_ref'](t).ravel()[0] # scalar
    x_tot = xref + e
    u_tot = uref + du[0]

    # total and reference dynamics
    f_tot = acrobot_dynamics(t, x_tot, u_tot, params['plant_params'])
    f_ref = acrobot_dynamics(t, xref, uref, params['plant_params'])

    return f_tot - f_ref    # error dynamics


acrobot_err = ctl.nlsys(
    acrobot_err_update, None, name="acrobot_err",
    inputs=('du',),
    outputs=('e1','e2','e3','e4'),
    states=('e1','e2','e3','e4'),
    params={
        'x_ref': x_d,
        'u_ref': u_d,
        'plant_params': plant_params
    }
)

# TODO: design MPC controller for trajectory tracking using acrobot_err system

def wrap_angle(angle):
    """Wrap angle to [-pi, pi)."""
    return (angle + np.pi) % (2*np.pi) - np.pi

# Thresholds for upright region, you can change them as needed
th1_thresh = 0.25   # rad (~14 deg) around pi
th2_thresh = 0.5    # rad (~29 deg) around 0
w1_thresh  = 1.0    # rad/s
w2_thresh  = 1.0    # rad/s

# function to check if state is in upright region
def in_upright_region(x):
    th1, th2, th1d, th2d = x
    e_th1 = wrap_angle(th1 - np.pi)
    e_th2 = wrap_angle(th2 - 0.0)

    cond_angle = (abs(e_th1) < th1_thresh) and (abs(e_th2) < th2_thresh) # check angles
    cond_vel   = (abs(th1d) < w1_thresh) and (abs(th2d) < w2_thresh) # check angular velocities

    return cond_angle and cond_vel

u_max = 5.0   # input limit for safety
mode  = "swingup"   # initial mode

def hybrid_controller(x, t):
    global mode

    if mode == "swingup":
        # 1) Use your trajectory-tracking controller (OCP/MPC/PD+FF)
        u = u_track(x, t)

        # 2) Check if we can switch to LQR
        if should_switch_to_lqr(x):
            mode = "stabilize"

    elif mode == "stabilize":
        # 1) Use LQR controller around upright
        u = u_lqr(x)

        # 2) Optionally: if we drift too far → go back to swing-up
        if should_switch_to_swingup(x):
            mode = "swingup"

    else:
        raise ValueError("Unknown mode: " + mode)

    # Saturate the torque for safety
    u = float(np.clip(u, -u_max, u_max))
    return u

# Create and inspect the Gymnasium Acrobot environment
env = gym.make("Acrobot-v1", render_mode=None)
obs, info = env.reset()
print("Observation space:", env.observation_space)
print("Action space:", env.action_space)
print("Initial observation:", obs)

env.close()

from stable_baselines3 import PPO
import numpy as np
from stable_baselines3.common.evaluation import evaluate_policy

# Use the wrapper
env = gym.make("Acrobot-v1")
model = PPO("MlpPolicy", env, verbose=1) # Create PPO model instance
model.learn(total_timesteps=200_000) # Train for 200,000 timesteps

# Save model
model.save("ppo_acrobot") # this will save the trained model to a zip file

# Evaluate
mean_reward, std_reward = evaluate_policy(model, env, n_eval_episodes=20)
print(f"Mean reward over 20 eval episodes: {mean_reward:.2f} ± {std_reward:.2f}")

env.close()

import gymnasium as gym
from stable_baselines3 import PPO
from stable_baselines3.common.vec_env import VecVideoRecorder, DummyVecEnv
from IPython.display import Video

# Create a vectorized environment for video recording
def make_env():
    return gym.make("Acrobot-v1", render_mode="rgb_array")

env = DummyVecEnv([make_env])

# Wrap with video recorder
env = VecVideoRecorder(
    env,
    "videos/",
    record_video_trigger=lambda x: x == 0,  # Record first episode
    video_length=500,  # Maximum video length
    name_prefix="acrobot_ppo"
)

# Load the trained model
model = PPO.load("ppo_acrobot")

# Run one episode with video recording
obs = env.reset()
for _ in range(500):
    action, _ = model.predict(obs, deterministic=True)
    obs, reward, done, info = env.step(action)
    if done[0]:
        break

env.close()

# Display the recorded video
Video("videos/acrobot_ppo-step-0-to-step-500.mp4", width=600)

import gymnasium as gym
from gymnasium import spaces
import numpy as np

# Action wrapper to convert discrete actions to continuous torque
class ContinuousAcrobotActionWrapper(gym.Wrapper):
    """
    Convert discrete Acrobot-v1 actions into continuous torque in [-torque_limit, torque_limit].
    The underlying env is still Acrobot-v1, but we override the torque used in the dynamics.
    """
    def __init__(self, env, torque_limit=5.0):
        super().__init__(env)
        self.torque_limit = float(torque_limit)
        # Continuous 1-D action: torque in [-torque_limit, torque_limit]
        self.action_space = spaces.Box(
            low=np.array([-self.torque_limit], dtype=np.float32),
            high=np.array([self.torque_limit], dtype=np.float32),
            dtype=np.float32,
        )
        # for reward wrapper
        self.unwrapped.last_u = 0.0

    def reset(self, **kwargs):
        obs, info = self.env.reset(**kwargs)
        self.unwrapped.last_u = 0.0
        return obs, info

    def step(self, action):
        # action is expected to be a 1D array or scalar; clip to [-torque_limit, torque_limit]
        if isinstance(action, np.ndarray):
            u = float(action.squeeze())
        else:
            u = float(action)
        u = np.clip(u, -self.torque_limit, self.torque_limit)
        self.unwrapped.last_u = u

        # Hack: temporarily set AVAIL_TORQUE so that super().step(1) uses our u
        orig_avail = self.env.unwrapped.AVAIL_TORQUE.copy()
        self.env.unwrapped.AVAIL_TORQUE[:] = [u, u, u]

        # Use middle index (1) – all entries are the same anyway
        obs, reward, terminated, truncated, info = self.env.step(1)

        # Restore original discrete torque set
        self.env.unwrapped.AVAIL_TORQUE[:] = orig_avail

        return obs, reward, terminated, truncated, info

# Reward wrapper to implement OCP-like reward
class AcrobotOCPRewardWrapper(gym.Wrapper):
    """
    Replace Acrobot reward by r = -u^T Q_u u + height reward + terminal bonus (on success).
    Assumes the wrapped env sets env.unwrapped.last_u to the actual torque used.
    """
    def __init__(self, env, Q_u=1.0, height_weight=5.0, terminal_bonus=200.0,
                 failure_penalty=0.0, success_height_threshold=1.8):
        super().__init__(env)
        self.Q_u = float(Q_u)
        self.height_weight = float(height_weight) # weight for height reward
        self.terminal_bonus = float(terminal_bonus) # bonus for success
        self.failure_penalty = float(failure_penalty) # penalty on failure (timeout)
        self.success_height_threshold = float(success_height_threshold)  # height threshold for success

    def step(self, action):
        obs, _, terminated, truncated, info = self.env.step(action)

        u = float(getattr(self.env.unwrapped, "last_u", 0.0))
        cos_th1, sin_th1, cos_th2, sin_th2, th1_dot, th2_dot = obs

        height = -cos_th1 - (cos_th1 * cos_th2 - sin_th1 * sin_th2)

        # Reward components
        height_reward = self.height_weight * height
        control_cost = self.Q_u * (u ** 2)
        velocity_penalty = 0.01 * (th1_dot**2 + th2_dot**2)
        reward = height_reward - control_cost - velocity_penalty

        # OVERRIDE termination logic: only terminate if reaching near-upright (height > 1.8)
        # or if episode times out (truncated)
        if height >= self.success_height_threshold:
            terminated = True  # True success
            reward += self.terminal_bonus
        else:
            terminated = False  # Keep going even if original env terminated at height > 1

        if truncated and self.failure_penalty != 0.0:
            reward -= self.failure_penalty

        return obs, reward, terminated, truncated, info

from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env

# Function to create the wrapped Acrobot environment
def make_acrobot_ocp_env():
    base_env = gym.make("Acrobot-v1")
    env = ContinuousAcrobotActionWrapper(base_env, torque_limit=5.0)
    env = AcrobotOCPRewardWrapper(
        env,
        Q_u=0.1,
        height_weight=5.0,
        terminal_bonus=200.0,
        success_height_threshold=1.7  # Require height for success, you may start with a smaller value so that it is easier to learn
    )
    return env

vec_env = make_vec_env(make_acrobot_ocp_env, n_envs=8) # Vectorized environment for parallel training

model = PPO("MlpPolicy", vec_env, verbose=1)
model.learn(total_timesteps=500_000)  # May need more timesteps
model.save("ppo_acrobot_ocp_continuous")

from stable_baselines3 import PPO
from stable_baselines3.common.vec_env import VecVideoRecorder, DummyVecEnv
from IPython.display import Video

# Create the wrapped environment with rgb_array render mode for video recording
def make_acrobot_ocp_env_visual():
    base_env = gym.make("Acrobot-v1", render_mode="rgb_array")
    env = ContinuousAcrobotActionWrapper(base_env, torque_limit=5.0)
    env = AcrobotOCPRewardWrapper(
        env,
        Q_u=0.1,
        height_weight=5.0,
        terminal_bonus=200.0,
        success_height_threshold=1.7
    )
    return env

# Create vectorized environment for video recording
env = DummyVecEnv([make_acrobot_ocp_env_visual])

# Wrap with video recorder
env = VecVideoRecorder(
    env,
    "videos/",
    record_video_trigger=lambda x: x == 0,  # Record first episode
    video_length=500,
    name_prefix="acrobot_ocp_continuous"
)

# Load the trained model
model = PPO.load("ppo_acrobot_ocp_continuous")

# Run one episode with video recording
obs = env.reset()
for _ in range(500):
    action, _ = model.predict(obs, deterministic=True)
    obs, reward, done, info = env.step(action)
    if done[0]:
        break

env.close()

# Display the recorded video
Video("videos/acrobot_ocp_continuous-step-0-to-step-500.mp4",embed=True, width=600)

from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env
from stable_baselines3.common.evaluation import evaluate_policy

class AcrobotSwingUpWrapper(gym.Wrapper):
    """
    Modified Acrobot wrapper for swing-up only (no balance).
    Key changes:
    1. Terminates when target height is reached
    2. Reward function focuses on height achievement
    3. No balance phase - episode ends at success
    """
    def __init__(self,
                 env,
                 height_weight=10.0,
                 control_cost=0.1,
                 success_height=1.8,
                 success_bonus=100.0,
                 downward_penalty=5.0):
        super().__init__(env)
        self.height_weight = height_weight
        self.control_cost = control_cost
        self.success_height = success_height
        self.success_bonus = success_bonus
        self.downward_penalty = downward_penalty

    def reset(self, **kwargs):
        obs, info = self.env.reset(**kwargs)
        return obs, info

    def _is_downward(self, th1):
        """Check if the acrobot is in a downward configuration"""
        e_th1_down = min(np.abs(th1), np.abs(th1 - 2*np.pi))
        return e_th1_down < 0.5

    def step(self, action):
        obs, _, terminated, truncated, info = self.env.step(action)

        # Extract state from observation
        cos_th1, sin_th1, cos_th2, sin_th2, th1_dot, th2_dot = obs

        # Recover angles
        th1 = np.arctan2(sin_th1, cos_th1)

        # Calculate tip height
        height = -cos_th1 - (cos_th1 * cos_th2 - sin_th1 * sin_th2)

        # Get control input
        u = float(getattr(self.env.unwrapped, "last_u", 0.0))

        # Swing-up reward function (no balance components)
        # 1. Height reward (main objective)
        height_reward = self.height_weight * height

        # 2. Downward penalty (discourage staying down)
        downward_cost = 0.0
        if self._is_downward(th1):
            downward_cost = self.downward_penalty

        # 3. Control cost (energy efficiency)
        control_penalty = self.control_cost * (u**2)

        # Total reward (simplified for swing-up only)
        reward = height_reward - downward_cost - control_penalty

        # TERMINATE upon reaching target height (swing-up success)
        if height >= self.success_height:
            terminated = True
            reward += self.success_bonus  # Large bonus for success

        return obs, reward, terminated, truncated, info


# Updated environment creation for swing-up only
def make_acrobot_swingup_env():
    base_env = gym.make("Acrobot-v1", max_episode_steps=500)  # Shorter episodes for swing-up
    env = ContinuousAcrobotActionWrapper(base_env, torque_limit=5.0)
    env = AcrobotSwingUpWrapper(
        env,
        height_weight=10.0,
        control_cost=0.1,
        success_height=1.8,  # Height threshold for success
        success_bonus=100.0,
        downward_penalty=5.0
    )
    return env

# Train the model for swing-up only
print("Training the swing-up only policy...")
vec_env_swingup = make_vec_env(make_acrobot_swingup_env, n_envs=32)
model_swingup = PPO("MlpPolicy", vec_env_swingup, verbose=1,
                    learning_rate=3e-4, n_steps=2048)
model_swingup.learn(total_timesteps=500_000)  # Less timesteps needed for simpler task
model_swingup.save("ppo_acrobot_swingup")
print("Training complete!")

# Evaluate the trained policy
env_eval = make_acrobot_swingup_env()
mean_reward, std_reward = evaluate_policy(model_swingup, env_eval, n_eval_episodes=10)
print(f"Mean reward over 10 episodes: {mean_reward:.2f} ± {std_reward:.2f}")
env_eval.close()

from stable_baselines3 import PPO
from stable_baselines3.common.vec_env import VecVideoRecorder, DummyVecEnv
import gymnasium as gym
from IPython.display import Video

# Create environment wrapper for video recording
def make_acrobot_balance_env_video():
    base_env = gym.make("Acrobot-v1", render_mode="rgb_array", max_episode_steps=1000)
    env = ContinuousAcrobotActionWrapper(base_env, torque_limit=5.0)
    env = AcrobotSwingUpWrapper(
        env,
        height_weight=10.0,
        control_cost=0.1,
        success_height=1.8,
        success_bonus=100.0,
        downward_penalty=5.0
    )
    return env

# Load the trained model
model = PPO.load("ppo_acrobot_swingup")

# Record video
env = DummyVecEnv([make_acrobot_balance_env_video])
env = VecVideoRecorder(env, "videos/",
                       record_video_trigger=lambda x: x == 0,
                       video_length=1000,
                       name_prefix="acrobot_balance")

obs = env.reset()
for _ in range(5000):
    action, _ = model.predict(obs, deterministic=True)
    obs, reward, done, info = env.step(action)
    if done[0]:
        break

env.close()

# Display the video
Video("videos/acrobot_swingup-step-0-to-step-1000.mp4",embed=True, width=600)

import matplotlib.pyplot as plt
from stable_baselines3 import PPO
import numpy as np

# Create environment without rendering
def make_acrobot_balance_env_eval():
    base_env = gym.make("Acrobot-v1", max_episode_steps=1000)
    env = ContinuousAcrobotActionWrapper(base_env, torque_limit=5.0)
    env = AcrobotSwingUpWrapper(
        env,
        height_weight=10.0,
        control_cost=0.1,
        success_height=1.8,
        success_bonus=100.0,
        downward_penalty=5.0
    )
    return env

model = PPO.load("ppo_acrobot_swingup")
env = make_acrobot_balance_env_eval()

# Run multiple episodes and collect data
num_episodes = 5
for ep in range(num_episodes):
    obs, info = env.reset()
    done = False
    step_count = 0

    # Storage for this episode
    observations = []
    actions = []
    rewards = []

    while not done and step_count < 1000:
        action, _ = model.predict(obs, deterministic=True)
        observations.append(obs)
        actions.append(action)

        obs, reward, terminated, truncated, info = env.step(action)
        rewards.append(reward)
        done = terminated or truncated
        step_count += 1

    # Convert to numpy arrays
    observations = np.array(observations)
    actions = np.array(actions)

    # Extract angles from observations
    cos_th1, sin_th1 = observations[:, 0], observations[:, 1]
    cos_th2, sin_th2 = observations[:, 2], observations[:, 3]
    th1 = np.arctan2(sin_th1, cos_th1)
    th2 = np.arctan2(sin_th2, cos_th2)

    # Calculate tip height
    height = -cos_th1 - (cos_th1 * cos_th2 - sin_th1 * sin_th2)

    # Plot results for this episode
    fig, axes = plt.subplots(2, 2, figsize=(14, 10))

    # Plot angles
    axes[0, 0].plot(th1, label='θ₁')
    axes[0, 0].plot(th2, label='θ₂')
    axes[0, 0].axhline(y=np.pi, color='r', linestyle='--', alpha=0.5, label='θ₁ target (π)')
    axes[0, 0].axhline(y=0, color='g', linestyle='--', alpha=0.5, label='θ₂ target (0)')
    axes[0, 0].set_xlabel('Step')
    axes[0, 0].set_ylabel('Angle (rad)')
    axes[0, 0].set_title('Joint Angles')
    axes[0, 0].legend()
    axes[0, 0].grid(True)

    # Plot tip height
    axes[0, 1].plot(height)
    axes[0, 1].axhline(y=2.0, color='r', linestyle='--', alpha=0.5, label='Max height (upright)')
    axes[0, 1].set_xlabel('Step')
    axes[0, 1].set_ylabel('Height')
    axes[0, 1].set_title('Tip Height Over Time')
    axes[0, 1].legend()
    axes[0, 1].grid(True)

    # Plot control actions
    axes[1, 0].plot(actions)
    axes[1, 0].axhline(y=5.0, color='r', linestyle='--', alpha=0.5, label='Torque limit')
    axes[1, 0].axhline(y=-5.0, color='r', linestyle='--', alpha=0.5)
    axes[1, 0].set_xlabel('Step')
    axes[1, 0].set_ylabel('Torque (Nm)')
    axes[1, 0].set_title('Control Input')
    axes[1, 0].legend()
    axes[1, 0].grid(True)

    # Plot rewards
    axes[1, 1].plot(rewards)
    axes[1, 1].set_xlabel('Step')
    axes[1, 1].set_ylabel('Reward')
    axes[1, 1].set_title('Reward per Step')
    axes[1, 1].grid(True)

    plt.suptitle(f'Episode {ep+1}: {step_count} steps, Final reward: {rewards[-1]:.2f}')
    plt.tight_layout()
    plt.show()

    print(f"Episode {ep+1}: Steps = {step_count}, Total reward = {sum(rewards):.2f}")

env.close()

#TODO: implement the missing parts