Lab 10: RL: Policy Gradient Methods

# @title
import gymnasium as gym
from gymnasium import spaces # spaces are used to define action and observation spaces
import numpy as np

# Define the set of possible actions and their visual representations
ACTIONS = [0, 1, 2, 3, 4]  # up, right, down, left, stay
ACTION_ARROWS = {0:"↑", 1:"→", 2:"↓", 3:"←", 4:"*"}


class GridWorld(gym.Env):

    """Custom GridWorld environment compatible with Gymnasium API.
    Note that Gymnasium environments are typically implemented as classes that inherit from gym.Env.
    This allows them to integrate seamlessly with the Gymnasium framework,
    which provides tools for reinforcement learning research and development.
    """

    def __init__(self,
                 grid_shape=(5, 5),
                 forbidden=None,
                 goal=None,
                 goal_reward=10.0,
                 step_reward=-0.01,
                 forbidden_reward=-5.0,
                 boundary_reward=-1.0,
                 render_mode=None):
        """
        Initializes the GridWorld environment.

        Args:
            grid_shape (tuple): The shape of the grid (rows, columns).
            forbidden (list of tuples): A list of coordinates for forbidden states.
            goal (tuple): The coordinates of the goal state.
            goal_reward (float): The reward for reaching the goal.
            step_reward (float): The reward for a regular step.
            forbidden_reward (float): The reward for entering a forbidden state.
            boundary_reward (float): The reward for hitting a boundary.
            render_mode (str, optional): The mode for rendering the environment.
        """
        super().__init__()

        self.grid_shape = grid_shape
        # Default locations for forbidden cells and the goal if not provided
        self.forbidden = forbidden if forbidden is not None else [(1,1), (1,2),(3,1), (3,3) ,(4,1)]
        # Example forbidden positions
        self.goal = goal if goal is not None else (3,2)
        # Example goal position

        # Reward structure
        self.goal_reward = goal_reward
        self.step_reward = step_reward
        self.forbidden_reward = forbidden_reward
        self.boundary_reward = boundary_reward

        # Discount factor
        self.render_mode = render_mode

        # Agent's current position, initialized in reset()
        self.agent_pos = None

        # Define Gymnasium spaces
        # The action space is discrete with a size equal to the number of actions.
        self.action_space = spaces.Discrete(len(ACTIONS))
        # The observation space represents the agent's position on the grid.
        self.observation_space = spaces.MultiDiscrete(list(grid_shape))


    def in_bounds(self, s):
        """Check if a state 's' is within the grid boundaries."""
        r, c = s
        R, C = self.grid_shape
        return 0 <= r < R and 0 <= c < C

    def is_forbidden(self, s):
        """Check if a state 's' is a forbidden state."""
        return s in self.forbidden

    def states(self):
        """Return a list of all possible states (coordinates) in the grid."""
        valid_states = []
        R, C = self.grid_shape
        for r in range(R):
            for c in range(C):
                valid_states.append((r,c))
        return valid_states

    def reset(self, seed=None, options=None):
        """
        Reset the environment to a RANDOM valid position.
        """
        super().reset(seed=seed)

        # Loop until we find a valid starting position
        while True:
            # Randomly pick a row and column
            r = np.random.randint(self.grid_shape[0])
            c = np.random.randint(self.grid_shape[1])
            s = (r, c)

            # Ensure we don't start in a forbidden cell or on the goal
            # if not self.is_forbidden(s) and s != self.goal:
            #     self.agent_pos = s
            #     break
            if s != self.goal:
                self.agent_pos = s
                break

        obs = np.array(self.agent_pos, dtype=np.int32)
        return obs, {}

    def step(self, action):
        """
        Execute one timestep in the environment.
        The agent takes an action, and the environment transitions to a new state.

        Args:
            action (int): The action to be taken by the agent.

        Returns:
            tuple: A tuple containing the new observation, the reward,
                   a boolean indicating if the episode is terminated,
                   a boolean indicating if the episode is truncated,
                   and an info dictionary.
        """
        assert self.action_space.contains(action), f"Invalid action {action}"

        r, c = self.agent_pos

        # Map the action to a new position (nr, nc for new row, new col)
        if action == 0:    nr, nc = r - 1, c     # up
        elif action == 1:  nr, nc = r, c + 1     # right
        elif action == 2:  nr, nc = r + 1, c     # down
        elif action == 3:  nr, nc = r, c - 1     # left
        else:              nr, nc = r, c         # stay

        # Check for environment boundaries
        terminated, truncated, info = False, False, {}
        next_pos = (nr, nc)
        if not self.in_bounds(next_pos):
            next_pos = self.agent_pos  # Agent stays in the same position
            reward = self.boundary_reward
        elif self.is_forbidden(next_pos):
            # next_pos = self.agent_pos  # Agent stays in the same position
            reward = self.forbidden_reward
        elif next_pos == self.goal:
            reward = self.goal_reward
            terminated = True
        elif action == 4:
            reward = -1
        else:
            reward = 0

        # Update the agent's position
        self.agent_pos = next_pos

        # Prepare the return values
        obs = np.array(self.agent_pos, dtype=np.int32)

        return obs, reward, terminated, truncated, info

    def render(self, P, V=None, ax=None, show=True):
        """
        Renders the greedy policy as arrows and overlays the value for each state in a matplotlib plot.
        This function will be used as a class method for the env class.

        Args:
            P (dict): The policy mapping states (r, c) to actions (int).
            V (dict): Optional. A dictionary mapping states to their values.
            ax: Optional matplotlib axes object. If provided, plots on this axes instead of creating a new figure.
            show (bool): Whether to call plt.show() at the end. Default True. Set False when creating subplots.

        Returns:
            tuple: (fig, ax) - The figure and axes objects.
        """
        import matplotlib.pyplot as plt
        import numpy as np

        R, C = self.grid_shape
        if ax is None:
            fig, ax = plt.subplots(figsize=(7, 6))
        else:
            fig = ax.figure
        ax.set_xlim(-0.5, C - 0.5)
        ax.set_ylim(-0.5, R - 0.5)
        ax.set_xticks(np.arange(-0.5, C, 1))
        ax.set_yticks(np.arange(-0.5, R, 1))
        ax.grid(True, color='gray', linewidth=1, linestyle='-')
        ax.set_aspect('equal')
        ax.invert_yaxis()  # Match array indexing

        # Draw cell highlights for forbidden/goal/start
        for r in range(R):
            for c in range(C):
                s = (r, c)
                if self.is_forbidden(s):
                    rect = plt.Rectangle((c-0.5, r-0.5), 1, 1, color='gray', alpha=0.5)
                    ax.add_patch(rect)
                    ax.text(c, r, "X", ha='center', va='center', color='black', fontsize=14)
                elif s == self.goal:
                    rect = plt.Rectangle((c-0.5, r-0.5), 1, 1, color='yellow', alpha=0.45)
                    ax.add_patch(rect)
                    ax.text(c, r, "G", ha='center', va='center', color='black', fontsize=14)

        # Print arrows and values in every cell (except forbidden)
        for r in range(R):
            for c in range(C):
                s = (r, c)
                # Print policy arrow
                if (r, c) in P:
                    arrow = ACTION_ARROWS[P[(r, c)]]
                else:
                    arrow = ""
                if (V is not None) and ((r, c) in V):
                    valstr = f"{V[(r, c)]:.2f}"
                else:
                    valstr = ""

                # Place arrow in cell
                ax.text(
                    c, r-0.18, arrow,
                    ha='center', va='center', color='darkblue', fontsize=22, fontweight='bold', zorder=10
                )

                ax.text(
                    c, r+0.22, valstr,
                    ha='center', va='center', color='crimson', fontsize=11, alpha=0.99,
                    zorder=10
                )

        # Label grid
        for x in range(C):
            ax.text(x, -0.75, str(x), ha='center', va='center', color='black')
        for y in range(R):
            ax.text(-0.75, y, str(y), ha='center', va='center', color='black')

        ax.set_title("Policy and Values")
        ax.tick_params(bottom=False, left=False, right=False, top=False,
                      labelbottom=False, labelleft=False, labeltop=False, labelright=False)
        if show:
            plt.tight_layout()
            plt.show()

        return fig, ax

import torch # PyTorch library for building and training neural networks
import torch.nn as nn # Neural network module from PyTorch
import torch.optim as optim # Optimization algorithms from PyTorch

# Create a function to convert state to one-hot:
def state_to_onehot(state, grid_shape=(5, 5)):
    """Convert (row, col) to one-hot vector"""
    r, c = state
    index = r * grid_shape[1] + c  # Flatten to single index
    onehot = torch.zeros(grid_shape[0] * grid_shape[1])
    onehot[index] = 1.0
    return onehot.unsqueeze(0)

LR = 0.0003 # Learning rate
GAMMA = 0.99 # Discount factor
EPISODES = 20000 # Number of training episodes
MAX_STEPS = 50 # Max steps per episode
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

#----------------------------------------------
# Initialize environment, policy network, and optimizer
#-----------------------
env = GridWorld()

state_dim = env.grid_shape[0] * env.grid_shape[1]  # state dimensions: 25 for 5x5 grid
action_dim = env.action_space.n # action dimensions: 5 actions

# Define policy network using nn.Sequential
# Input: 25-dimensional one-hot vector
# Output: 5 logits (one per action)
# Architecture: Linear(25 → 128) → ReLU → Linear(128 → 5)
policy_net = nn.Sequential(
    nn.Linear(state_dim, 128),
    nn.ReLU(),
    nn.Linear(128, action_dim)
).to(device)

# Define optimizer which is used to update the policy network parameters
optimizer = optim.Adam(policy_net.parameters(), lr=LR)
return_list = []

#----------------------------------------------
# Training loop
#----------------------------------------------
print("Training started...")
for episode in range(EPISODES):
    log_probs = []
    rewards = []

    state, info = env.reset()
    episode_return = 0

    # Generate an episode
    for t in range(MAX_STEPS):
        state_tensor = state_to_onehot(state, env.grid_shape).to(device)

        # Get action logits from policy network
        # logits = [f_θ(s,0), f_θ(s,1), ..., f_θ(s,4)] are raw scores for each action
        logits = policy_net(state_tensor)

        # Categorical(logits=logits) applies softmax internally:
        # π_θ(a|s) = exp(f_θ(s,a)) / Σ_a' exp(f_θ(s,a'))
        # This converts logits into a valid probability distribution over actions
        action_dist = torch.distributions.Categorical(logits=logits)

        # Sample an action from the distribution π_θ(a|s)
        action = action_dist.sample()

        # Compute log π_θ(a_t|s_t) for the sampled action (needed for policy gradient)
        log_prob = action_dist.log_prob(action)

        # Take the action in the environment
        next_state, reward, terminated, truncated, _ = env.step(action.item())

        # Store log probability and reward
        log_probs.append(log_prob)
        rewards.append(reward)

        # Move to the next state
        state = next_state
        # Accumulate episode return
        episode_return += reward

        if terminated or truncated:
            break

    # Compute action values (returns) G_t for each time step t
    returns = []
    G = 0
    for r in reversed(rewards):
        G = r + GAMMA * G
        returns.insert(0, G)
    # Convert returns to tensor
    returns = torch.tensor(returns, dtype=torch.float32).to(device)

    # Compute loss -Σ log π(a_t|s_t) * G_t
    loss = 0
    for log_prob, G_t in zip(log_probs, returns):
        loss += -log_prob * G_t # loss is a tensor in PyTorch since log_prob is a tensor

    # Update policy (i.e., update the parameters of the policy network)
    optimizer.zero_grad() # Clear previous gradients
    loss.backward() # Backpropagate to compute gradients (since loss is a tensor, it has a .backward() method)
    optimizer.step() # Update policy network parameters

    return_list.append(episode_return)

    if (episode+1) % 1000 == 0:
        print(f"Episode {episode+1}/{EPISODES}, Return: {episode_return:.2f}")

# Visualization
policy_dict = {}
with torch.no_grad():
    for r in range(env.grid_shape[0]):
        for c in range(env.grid_shape[1]):
            s = (r, c)
            if s == env.goal:
                continue

            s_tensor = state_to_onehot(s, env.grid_shape).to(device)

            logits = policy_net(s_tensor)
            best_action = torch.argmax(logits).item()
            policy_dict[s] = best_action

env.render(policy_dict)

try:
    import stable_baselines3
except ImportError:
    !pip install stable_baselines3
    import stable_baselines3

# With the SB3 library and a Gymnasium-compatible environment, training can be done in just a few lines of code.
from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env

# 1. create the environment
env = GridWorld()


# 2. Initialize PPO model
policy_kwargs = dict(
    net_arch=[16]  # Single hidden layer with 16 neurons
)
model = PPO("MlpPolicy",  # Multi-layer perceptron policy
            env,  # the environment
            verbose=1,  # print out info during training
            gamma=0.99,  # discount factor, if not provided, default is 0.99
            learning_rate=1e-2,  # learning rate, if not provided, default is 3e-4
            policy_kwargs=policy_kwargs # policy network architecture: if this is not provided, default is [64, 64]
            )


# 3. Train model
model.learn(total_timesteps=50_000)
print("training finished")
model.save("ppo_gridworld")

## Now let us render the learned policy

# --- 1. Prepare to extract the policy ---
policy_dict = {}
R, C = env.grid_shape

# --- 2. Iterate over all states and find the best action ---
for r in range(R):
    for c in range(C):
        s = (r, c)

        # Skip the goal state, as no action is taken there
        if s == env.goal:
            continue

        # The input to the model's predict() must be a NumPy array.
        state_np = np.array([r, c], dtype=np.int32)

        # model.predict returns the action (0-4) and the hidden state (which we ignore)
        best_action, _ = model.predict(state_np, deterministic=True)

        # Store the greedy action for the state
        policy_dict[s] = int(best_action)

# --- 3. Render the policy ---
env.render(policy_dict)

import imageio # for creating videos
import gymnasium as gym # Gymnasium library for reinforcement learning environments
from stable_baselines3 import PPO # PPO algorithm
from stable_baselines3.common.env_util import make_vec_env # for creating vectorized environments
import numpy as np
from IPython.display import Video # for displaying videos in Jupyter notebooks

# 1. Create training env
# Pendulum-v1 comes from Gymnasium's built-in environments
env = make_vec_env("Pendulum-v1", n_envs=4)

# 2. Train PPO model
model = PPO(
    "MlpPolicy",
    env,
    verbose=1,
    learning_rate=3e-4,
    gamma=0.99,
)

model.learn(total_timesteps=100_000) # the number of training timesteps can be adjusted

# 3. Save the trained model
model.save("ppo_pendulum_v1")
print("Training finished and model saved.")

# Create rendering env
render_env = gym.make(
    "Pendulum-v1",
    render_mode="rgb_array", # to enable image rendering
    max_episode_steps=3000
)

# Load the trained model in the previous cell
model = PPO.load("ppo_pendulum_v1")

# Record video using the imageio library
writer = imageio.get_writer("pendulum_video.mp4", fps=30)

# Reset the environment
obs, info = render_env.reset()

# Run the trained model and record frames
for _ in range(300): # generate 300 frames
    action, _ = model.predict(obs, deterministic=True) # get action from the trained model
    obs, reward, terminated, truncated, info = render_env.step(action) # take action in env

    frame = render_env.render()
    writer.append_data(frame)

    if terminated or truncated:
        break

writer.close()
render_env.close()
Video("pendulum_video.mp4", embed=True, width=600)

import imageio
import gymnasium as gym
from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env
import numpy as np
from IPython.display import Video


# 1. Create training env
env = make_vec_env("CartPole-v1", n_envs=1) #n_envs=1 means no vectorization

# 2. Train PPO model
model = PPO(
    "MlpPolicy",
    env,
    verbose=1,
    learning_rate=3e-4,
    gamma=0.99,
)

model.learn(total_timesteps=300_000)

# Save model
model.save("ppo_CartPole_v1")
print("Training finished and model saved.")

# 3. Create rendering env
render_env = gym.make(
    "CartPole-v1",
    render_mode="rgb_array",
    max_episode_steps=2000
)

# Load model
model = PPO.load("ppo_CartPole_v1")

# 4. Record video
writer = imageio.get_writer("pendulum_video.mp4", fps=30)

obs, info = render_env.reset()

for _ in range(300):
    action, _ = model.predict(obs, deterministic=True)
    obs, reward, terminated, truncated, info = render_env.step(action)

    frame = render_env.render()
    writer.append_data(frame)

    if terminated or truncated:
        break

writer.close()
render_env.close()
Video("pendulum_video.mp4", embed=True, width=600)

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


class ContinuousCartPoleWrapper(gym.Wrapper):
    """
    Wrapper that turns CartPole-v1 into a continuous-control,
    LQR-style stabilization task.

    - Action: u in [-u_max, u_max]  (continuous force)
    - State:  [x, x_dot, theta, theta_dot]
    - Reward: r = -(x^T Q x + u^T R u) * dt, with optional terminal penalty.
    """

    def __init__(
        self,
        env: gym.Env,
        u_max: float = 10.0,
        Q: np.ndarray | None = None,
        R: float | None = None,
        scale_with_dt: bool = True,
        terminal_cost: float = 100.0,
    ):
        super().__init__(env)

        # Continuous 1D action (force)
        self.u_max = float(u_max)
        self.action_space = spaces.Box(
            low=np.array([-self.u_max], dtype=np.float32),
            high=np.array([self.u_max], dtype=np.float32),
            dtype=np.float32,
        )

        # Keep the original observation space (4-dim Box)
        self.observation_space = env.observation_space

        # Default LQR weights if not provided
        # Emphasize pole angle θ and cart position x
        if Q is None:
            # state = [x, x_dot, theta, theta_dot]
            Q = np.diag([1.0, 0.1, 10.0, 0.1])
        if R is None:
            R = 0.01  # scalar for u^2

        self.Q = np.array(Q, dtype=np.float64)
        self.R = float(R)

        self.scale_with_dt = scale_with_dt
        self.terminal_cost = float(terminal_cost)

    # Just delegate reset to inner env and return same observation
    def reset(self, **kwargs):
        obs, info = self.env.reset(**kwargs)
        return obs, info

    def step(self, action):
        # ---- 1) Continuous control input u ----
        if isinstance(action, np.ndarray):
            u = float(action.squeeze())
        else:
            u = float(action)
        u = np.clip(u, -self.u_max, self.u_max)

        # Use the *unwrapped* CartPole to access physics params and state
        env = self.env.unwrapped

        # ---- 2) Current state ----
        x, x_dot, theta, theta_dot = env.state

        # ---- 3) Dynamics from Gymnasium's CartPole (but with continuous u) ----
        gravity = env.gravity
        masscart = env.masscart
        masspole = env.masspole
        total_mass = masscart + masspole
        length = env.length          # actually half the pole length
        polemass_length = masspole * length
        tau = env.tau                # time step

        costheta = np.cos(theta)
        sintheta = np.sin(theta)

        # This is exactly CartPole's internal derivation, but with "force = u"
        temp = (u + polemass_length * theta_dot**2 * sintheta) / total_mass
        theta_acc = (gravity * sintheta - costheta * temp) / (
            length * (4.0 / 3.0 - masspole * costheta**2 / total_mass)
        )
        x_acc = temp - polemass_length * theta_acc * costheta / total_mass

        # Euler integration
        x = x + tau * x_dot
        x_dot = x_dot + tau * x_acc
        theta = theta + tau * theta_dot
        theta_dot = theta_dot + tau * theta_acc

        env.state = (x, x_dot, theta, theta_dot)

        # ---- 4) Termination condition (same as CartPole) ----
        terminated = (
            x < -env.x_threshold
            or x > env.x_threshold
            or theta < -env.theta_threshold_radians
            or theta > env.theta_threshold_radians
        )
        # We don't enforce a time limit here; TimeLimit wrapper can be applied outside if you like
        truncated = False

        """# ---- 5) LQR-like reward ----
        state_vec = np.array([x, x_dot, theta, theta_dot], dtype=np.float64)
        state_cost = state_vec @ self.Q @ state_vec
        control_cost = self.R * (u ** 2)

        cost = state_cost + control_cost
        if self.scale_with_dt:
            cost *= tau
        reward = -float(cost)
        """
        # ---- LQR-like reward with boundary shaping ----
        state_vec = np.array([x, x_dot, theta, theta_dot], dtype=np.float64)
        state_cost = state_vec @ self.Q @ state_vec
        control_cost = self.R * (u ** 2)

        # Soft barrier on cart position near limits
        x_norm = abs(x) / env.x_threshold  # 0 at center, 1 at boundary
        boundary_cost = 5.0 * (x_norm ** 4)  # smooth, grows steeply near boundary

        # Extra angle robustness term (helps with large initial angles)
        angle_cost = 0.5 * (1.0 - np.cos(theta)) * 20.0  # ~10 when upside down

        cost = state_cost + control_cost + boundary_cost + angle_cost
        if self.scale_with_dt:
            cost *= tau

        reward = -float(cost) # total reward is negative cost


        # Add terminal cost when it fails
        if terminated:
            reward -= self.terminal_cost

        # Add terminal cost when it fails
        if terminated:
            reward -= self.terminal_cost

        obs = state_vec.astype(np.float32)
        info = {
            "u": u,
            "state_vec": obs,
            "state_cost": state_cost,
            "control_cost": control_cost,
            "instant_cost": cost,
        }

        return obs, reward, terminated, truncated, info

    # For rendering, just delegate
    def render(self):
        return self.env.render()

import gymnasium as gym
from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env

# Create the environment with continuous action wrapper
def make_env():
    base_env = gym.make("CartPole-v1")  # standard env
    env = ContinuousCartPoleWrapper(base_env, u_max=20.0)
    return env

env = make_vec_env(make_env, n_envs=1)

model = PPO(
    "MlpPolicy",
    env,
    verbose=1,
    learning_rate=3e-4,
    gamma=0.99
)

model.learn(total_timesteps=150_000)
model.save("ppo_continuous_cartpole")

import imageio
from IPython.display import Video

render_env = gym.make(
    "CartPole-v1",
    render_mode="rgb_array",  # important for frame capture
)
render_env = ContinuousCartPoleWrapper(render_env, u_max=20.0)

model = PPO.load("ppo_continuous_cartpole")

writer = imageio.get_writer("CartPole_LQR.mp4", fps=30)

obs, info = render_env.reset()
for _ in range(300):
    action, _ = model.predict(obs, deterministic=True)
    obs, reward, terminated, truncated, info = render_env.step(action)

    frame = render_env.render()
    writer.append_data(frame)

    if terminated or truncated:
        break

writer.close()
render_env.close()

Video("CartPole_LQR.mp4", embed=True, width=600)

class WideInitWrapper(gym.Wrapper):
    """
    Reset with wider ranges so PPO sees large initial angles.
    Ranges are tunable.
    """
    def __init__(
        self,
        env,
        x_range=1, # range for cart position
        xdot_range=1, # range for cart velocity
        theta_range=np.pi/10, # range for pole angle (up to ±18°)
        thetadot_range=2.0, # range for pole angular velocity
    ):
        super().__init__(env)
        self.x_range = x_range
        self.xdot_range = xdot_range
        self.theta_range = theta_range
        self.thetadot_range = thetadot_range

    def reset(self, **kwargs):
        obs, info = self.env.reset(**kwargs)
        rng = self.env.unwrapped.np_random

        x = rng.uniform(-self.x_range, self.x_range)
        x_dot = rng.uniform(-self.xdot_range, self.xdot_range)
        theta = rng.uniform(-self.theta_range, self.theta_range)
        theta_dot = rng.uniform(-self.thetadot_range, self.thetadot_range)

        self.env.unwrapped.state = (x, x_dot, theta, theta_dot)
        obs = np.array([x, x_dot, theta, theta_dot], dtype=np.float32)
        return obs, info

import gymnasium as gym
from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env

# Create the environment with continuous action wrapper
def make_env():
    base = gym.make("CartPole-v1")
    base = ContinuousCartPoleWrapper(base, u_max=20.0)  # or ContinuousLQRCartPoleWrapper
    base = WideInitWrapper(base, theta_range=np.pi/10, thetadot_range=1.0)
    return base

env = make_vec_env(make_env, n_envs=4)

model = PPO(
    "MlpPolicy",
    env,
    verbose=1,
    learning_rate=3e-4,
    gamma=0.99
)

model.learn(total_timesteps=500_000)
model.save("ppo_continuous_cartpole")

# Now render the trained model with wide init
import imageio
from IPython.display import Video

render_env = gym.make(
    "CartPole-v1",
    render_mode="rgb_array",  # important for frame capture
)
render_env = ContinuousCartPoleWrapper(render_env, u_max=20.0)
render_env = WideInitWrapper(render_env, theta_range=np.pi/10, thetadot_range=1.0)  # match training init

model = PPO.load("ppo_continuous_cartpole")

writer = imageio.get_writer("CartPole.mp4", fps=30)

obs, info = render_env.reset()
for _ in range(300):
    action, _ = model.predict(obs, deterministic=True)
    obs, reward, terminated, truncated, info = render_env.step(action)

    frame = render_env.render()
    writer.append_data(frame)

    if terminated or truncated:
        break

writer.close()
render_env.close()

Video("CartPole_LQR.mp4", embed=True, width=600)

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


class ContinuousSwingUpCartPoleWrapper(gym.Wrapper):
    """
    Continuous-action CartPole with swing-up style reward.

    - Action: u in [-u_max, u_max]
    - State:  [x, x_dot, theta, theta_dot]
    - Reward: encourage upright (cos(theta)) and penalize big x, velocities, and control.
    - Done:
        * success when near-upright for several consecutive steps
        * failure if cart goes out of bounds
        * episode time limit is handled by an outer TimeLimit wrapper.
    """

    def __init__(
        self,
        env: gym.Env,
        u_max: float = 10.0,  # max force
        x_fail_threshold: float = 2.4,  # cart position failure threshold
        upright_theta_thresh: float = np.deg2rad(12.0),  # setpoint for "upright"
        upright_theta_dot_thresh: float = 1.0,  # velocity threshold for "upright"
        upright_x_thresh: float = 0.5,  # position threshold for "upright"
        success_steps: int = 2,  # number of consecutive upright steps for success
        success_bonus: float = 50.0,  # reward for achieving upright
        failure_penalty: float = 50.0,  # penalty for failure
    ):
        super().__init__(env)

        # Continuous 1D action space
        self.u_max = float(u_max)
        self.action_space = spaces.Box(
            low=np.array([-self.u_max], dtype=np.float32),
            high=np.array([self.u_max], dtype=np.float32),
            dtype=np.float32,
        )

        # Same observation space as CartPole
        self.observation_space = env.observation_space

        # Swing-up termination thresholds
        self.x_fail_threshold = float(x_fail_threshold)
        self.upright_theta_thresh = float(upright_theta_thresh)
        self.upright_theta_dot_thresh = float(upright_theta_dot_thresh)
        self.upright_x_thresh = float(upright_x_thresh)

        # Success condition (must hold upright for N consecutive steps)
        self.success_steps = int(success_steps)
        self.steps_upright = 0  # counter

        self.success_bonus = float(success_bonus)
        self.failure_penalty = float(failure_penalty)

    def reset(self, **kwargs):
        # Let another wrapper handle initial conditions if present
        obs, info = self.env.reset(**kwargs)
        self.steps_upright = 0
        return obs, info

    def step(self, action):
        # ---- 1) Continuous control input u ----
        if isinstance(action, np.ndarray):
            u = float(action.squeeze())
        else:
            u = float(action)
        u = np.clip(u, -self.u_max, self.u_max)

        env = self.env.unwrapped

        # ---- 2) Current state ----
        x, x_dot, theta, theta_dot = env.state

        # ---- 3) Dynamics (same as CartPole, but with force = u) ----
        gravity = env.gravity
        masscart = env.masscart
        masspole = env.masspole
        total_mass = masscart + masspole
        length = env.length          # actually half pole length
        polemass_length = masspole * length
        tau = env.tau                # time step

        costheta = np.cos(theta)
        sintheta = np.sin(theta)

        temp = (u + polemass_length * theta_dot**2 * sintheta) / total_mass
        theta_acc = (gravity * sintheta - costheta * temp) / (
            length * (4.0 / 3.0 - masspole * costheta**2 / total_mass)
        )
        x_acc = temp - polemass_length * theta_acc * costheta / total_mass

        # Euler integration
        x = x + tau * x_dot
        x_dot = x_dot + tau * x_acc
        theta = theta + tau * theta_dot
        theta_dot = theta_dot + tau * theta_acc

        env.state = (x, x_dot, theta, theta_dot)

        # ---- 4) Termination: success/failure ----
        # "Upright" region: near upright and reasonably calm
        upright = (
            abs(theta) < self.upright_theta_thresh
            and abs(theta_dot) < self.upright_theta_dot_thresh
            and abs(x) < self.upright_x_thresh
        )

        if upright:
            self.steps_upright += 1
        else:
            self.steps_upright = 0

        # Success if we've stayed upright long enough
        success = self.steps_upright >= self.success_steps

        # Failure: cart leaves the track
        failure = abs(x) > self.x_fail_threshold

        terminated = success or failure
        truncated = False  # TimeLimit wrapper outside can handle max steps

        # ---- 5) Swing-up reward ----
        # Main term: cos(theta) is highest at upright (1) and lowest at downward (-1)
        r_angle = np.cos(theta)

        # Penalize cart displacement
        r_x = -0.01 * (x**2)

        # Stronger penalty on angular velocity
        w_xdot = 0.01
        w_thetadot = 0.03
        r_vel = -(w_xdot * x_dot**2 + w_thetadot * theta_dot**2)

        # Penalize control effort
        r_u = -0.001 * (u**2)

        # Penalize approaching boundaries
        r_boundary = -5.0 * (abs(x) / self.env.unwrapped.x_threshold) ** 4

        reward = float(r_angle + r_x + r_vel + r_u + r_boundary)

        # Terminal bonuses/penalties
        if success:
            reward += self.success_bonus
        if failure:
            reward -= self.failure_penalty

        obs = np.array([x, x_dot, theta, theta_dot], dtype=np.float32)
        info = {
            "u": u,
            "upright": upright,
            "success": success,
            "failure": failure,
            "steps_upright": self.steps_upright,
        }

        return obs, reward, terminated, truncated, info

    def render(self):
        return self.env.render()

class SwingUpInitWrapper(gym.Wrapper):
    """
    Override reset() so the pole starts near the downward position (theta ~ pi).
    """

    # define the angle noise in degrees, default is 10 degrees
    def __init__(self, env, angle_noise_deg: float = 10.0):
        super().__init__(env)
        self.angle_noise = np.deg2rad(angle_noise_deg)

    def reset(self, **kwargs):
        # Call base reset to set up RNG etc.
        obs, info = self.env.reset(**kwargs)

        rng = self.env.unwrapped.np_random

        x = 0.0
        x_dot = 0.0
        # Set theta near pi (downward) with some noise
        theta = np.pi + rng.uniform(-self.angle_noise, self.angle_noise)
        theta_dot = 0.0

        self.env.unwrapped.state = (x, x_dot, theta, theta_dot)

        obs = np.array([x, x_dot, theta, theta_dot], dtype=np.float32)
        return obs, info

import gymnasium as gym
from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env
from stable_baselines3.common.monitor import Monitor
import matplotlib.pyplot as plt


# 1) Build vectorized environment
def make_env():
    base_env = gym.make("CartPole-v1")
    base_env = Monitor(base_env, log_dir)   # <-- log episodic stats
    env = ContinuousSwingUpCartPoleWrapper(base_env, u_max=20.0)
    env = SwingUpInitWrapper(env, angle_noise_deg=15.0)
    return env

env = make_vec_env(make_env, n_envs=4)  # multiple envs helps PPO

# 2) Train PPO
model = PPO(
    "MlpPolicy",
    env,
    verbose=1,
    learning_rate=3e-4,
    gamma=0.99,
    n_steps=1024,
    batch_size=64,
)

model.learn(total_timesteps=1_000_000)
model.save("ppo_cartpole_swingup")
env.close()

import imageio
from stable_baselines3 import PPO
from IPython.display import Video

# 1) Rebuild render env with same wrappers
render_env = gym.make(
    "CartPole-v1",
    render_mode="rgb_array",
    max_episode_steps=5000,   # allow enough time for swing-up
)

render_env = ContinuousSwingUpCartPoleWrapper(render_env, u_max=20.0)
render_env = SwingUpInitWrapper(render_env, angle_noise_deg=15.0)

# 2) Load trained model
model = PPO.load("ppo_cartpole_swingup")

# 3) Record video
writer = imageio.get_writer("CartPole_swingup.mp4", fps=30)

obs, info = render_env.reset()

for _ in range(5000):
    action, _ = model.predict(obs, deterministic=True)
    obs, reward, terminated, truncated, info = render_env.step(action)

    frame = render_env.render()
    writer.append_data(frame)

    if terminated or truncated:
        break

writer.close()
render_env.close()

Video("CartPole_swingup.mp4", embed=True, width=600)

import numpy as np
import gymnasium as gym
import imageio
from stable_baselines3 import PPO
from IPython.display import Video


# ---------- 1. Upright condition used for switching ----------

# Tunable thresholds for the *handoff* to stabilizer
THETA_THRESH = np.deg2rad(3.0)   # ~3 degrees
THETA_DOT_THRESH = 1            # rad/s
X_THRESH = 0.7                    # meters


def is_upright(obs) -> bool:
    """
    Check if the current state is 'upright enough' to hand off
    from the swing-up policy to the stabilizer policy.

    obs = [x, x_dot, theta, theta_dot]
    """
    x, x_dot, theta, theta_dot = obs
    return (
        abs(theta) < THETA_THRESH
        and abs(theta_dot) < THETA_DOT_THRESH
        and abs(x) < X_THRESH
    )


# ---------- 2. Load trained models ----------
print("Loading models...")
swingup_model = PPO.load("ppo_cartpole_swingup")
stabilizer_model = PPO.load("ppo_continuous_cartpole")
print("Models loaded successfully")


# ---------- 3. Create evaluation environment ----------
render_env = gym.make(
    "CartPole-v1",
    render_mode="rgb_array",
    max_episode_steps=1000,
)

# Wrap with swing-up wrapper for physics
render_env = ContinuousSwingUpCartPoleWrapper(
    render_env,
    u_max=20.0,
    x_fail_threshold=2.4,
    success_steps=10**9,   # disable automatic success termination
)

# Start from downward position
render_env = SwingUpInitWrapper(render_env, angle_noise_deg=5.0)


# ---------- 4. Run combined episode and record video ----------
writer = imageio.get_writer("combined_cartpole.mp4", fps=30)

obs, info = render_env.reset()
mode = "swingup"
print(f"Starting episode in mode: {mode}")
print(f"Initial state: x={obs[0]:.3f}, theta={obs[2]:.3f} rad ({np.rad2deg(obs[2]):.1f} deg)")

for t in range(1000):
    # Check if we should switch to stabilizer
    if mode == "swingup" and is_upright(obs):
        mode = "stabilize"
        print(f"[step {t}] Switching to STABILIZER policy")
        print(f"  State: x={obs[0]:.3f}, theta={obs[2]:.3f} rad ({np.rad2deg(obs[2]):.1f} deg)")

    # Choose action based on current mode
    if mode == "swingup":
        action, _ = swingup_model.predict(obs, deterministic=True)
    else:
        action, _ = stabilizer_model.predict(obs, deterministic=True)

    obs, reward, terminated, truncated, info = render_env.step(action)

    # Record frame
    frame = render_env.render()
    writer.append_data(frame)

    if terminated or truncated:
        print(f"Episode ended at step {t}")
        print(f"  terminated={terminated}, truncated={truncated}")
        print(f"  Final state: x={obs[0]:.3f}, theta={obs[2]:.3f} rad")
        break

writer.close()
render_env.close()

print("Video saved to combined_cartpole.mp4")
Video("combined_cartpole.mp4", embed=True, width=600)