Optimizing Chip Layout with Reinforcement Learning for AMD

# chip_env.py
import gym
from gym import spaces
import numpy as np
import networkx as nx
from typing import List, Tuple, Dict
 
class ChipFloorplanEnv(gym.Env):
    """
    Custom OpenAI Gym environment for chip floorplanning.
    
    State: Flattened representation of block positions, netlist connectivity
    Action: (block_id, x_position, y_position)
    Reward: Multi-objective function balancing wire length, congestion, timing
    """
    
    def __init__(
        self,
        canvas_width: int = 100,
        canvas_height: int = 100,
        num_blocks: int = 20,
        netlist_path: str = None
    ):
        super(ChipFloorplanEnv, self).__init__()
        
        self.canvas_width = canvas_width
        self.canvas_height = canvas_height
        self.num_blocks = num_blocks
        
        # Load netlist (connectivity graph)
        self.netlist = self._load_netlist(netlist_path)
        
        # Block properties (width, height, power density)
        self.blocks = self._initialize_blocks()
        
        # Action space: (block_id, x, y)
        self.action_space = spaces.MultiDiscrete([
            num_blocks,  # Which block to place
            canvas_width,  # X position
            canvas_height  # Y position
        ])
        
        # State space: block positions + connectivity matrix
        state_dim = num_blocks * 2 + num_blocks * num_blocks
        self.observation_space = spaces.Box(
            low=0,
            high=max(canvas_width, canvas_height),
            shape=(state_dim,),
            dtype=np.float32
        )
        
        # Placement tracking
        self.placed_blocks = set()
        self.block_positions = {}
        self.placement_grid = np.zeros((canvas_width, canvas_height))
        
        # Metrics
        self.total_wire_length = 0
        self.max_congestion = 0
        self.timing_violations = 0
        self.thermal_hotspots = 0
    
    def _load_netlist(self, path: str) -> nx.Graph:
        """Load netlist from file and create connectivity graph"""
        G = nx.Graph()
        
        if path:
            with open(path, 'r') as f:
                for line in f:
                    # Parse netlist format: block1 block2 wire_weight
                    block1, block2, weight = line.strip().split()
                    G.add_edge(int(block1), int(block2), weight=float(weight))
        else:
            # Generate random connectivity for testing
            for i in range(self.num_blocks):
                for j in range(i + 1, self.num_blocks):
                    if np.random.rand() < 0.3:  # 30% connectivity
                        G.add_edge(i, j, weight=np.random.uniform(1, 10))
        
        return G
    
    def _initialize_blocks(self) -> Dict:
        """Initialize block properties (dimensions, power)"""
        blocks = {}
        for i in range(self.num_blocks):
            blocks[i] = {
                'width': np.random.randint(5, 15),
                'height': np.random.randint(5, 15),
                'power_density': np.random.uniform(0.1, 2.0),  # W/mm²
                'timing_criticality': np.random.uniform(0, 1)  # 0 = non-critical, 1 = critical
            }
        return blocks
    
    def reset(self) -> np.ndarray:
        """Reset environment to initial state"""
        self.placed_blocks = set()
        self.block_positions = {}
        self.placement_grid = np.zeros((self.canvas_width, self.canvas_height))
        
        self.total_wire_length = 0
        self.max_congestion = 0
        self.timing_violations = 0
        self.thermal_hotspots = 0
        
        return self._get_state()
    
    def _get_state(self) -> np.ndarray:
        """
        Construct state vector:
        - Block positions (x, y for each block)
        - Connectivity matrix (flattened)
        """
        # Block positions (0 if not placed yet)
        positions = []
        for i in range(self.num_blocks):
            if i in self.block_positions:
                positions.extend([
                    self.block_positions[i][0] / self.canvas_width,
                    self.block_positions[i][1] / self.canvas_height
                ])
            else:
                positions.extend([0, 0])
        
        # Connectivity matrix
        conn_matrix = nx.to_numpy_array(self.netlist, weight='weight')
        conn_flat = conn_matrix.flatten()
        
        state = np.concatenate([positions, conn_flat])
        return state.astype(np.float32)
    
    def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, dict]:
        """
        Execute placement action and return new state, reward, done flag
        
        Args:
            action: [block_id, x, y]
        
        Returns:
            state, reward, done, info
        """
        block_id, x, y = action
        
        # Check if block already placed
        if block_id in self.placed_blocks:
            return self._get_state(), -10, False, {'error': 'block_already_placed'}
        
        # Check if placement is valid (no overlap)
        if not self._is_valid_placement(block_id, x, y):
            return self._get_state(), -10, False, {'error': 'invalid_placement'}
        
        # Place block
        self._place_block(block_id, x, y)
        
        # Calculate reward
        reward = self._calculate_reward()
        
        # Check if all blocks placed
        done = len(self.placed_blocks) == self.num_blocks
        
        # Additional metrics
        info = {
            'wire_length': self.total_wire_length,
            'congestion': self.max_congestion,
            'timing_violations': self.timing_violations,
            'thermal_hotspots': self.thermal_hotspots,
            'blocks_placed': len(self.placed_blocks)
        }
        
        return self._get_state(), reward, done, info
    
    def _is_valid_placement(self, block_id: int, x: int, y: int) -> bool:
        """Check if block can be placed at (x, y) without overlap"""
        block = self.blocks[block_id]
        
        # Check bounds
        if x + block['width'] > self.canvas_width:
            return False
        if y + block['height'] > self.canvas_height:
            return False
        
        # Check overlap with existing blocks
        for i in range(x, x + block['width']):
            for j in range(y, y + block['height']):
                if self.placement_grid[i, j] > 0:
                    return False
        
        return True
    
    def _place_block(self, block_id: int, x: int, y: int):
        """Place block at (x, y) and update grid"""
        block = self.blocks[block_id]
        
        # Mark grid cells as occupied
        for i in range(x, x + block['width']):
            for j in range(y, y + block['height']):
                self.placement_grid[i, j] = block_id + 1  # +1 to avoid 0
        
        # Store position
        self.block_positions[block_id] = (x, y)
        self.placed_blocks.add(block_id)
    
    def _calculate_reward(self) -> float:
        """
        Multi-objective reward function:
        - Minimize wire length (HPWL - Half-Perimeter Wire Length)
        - Minimize congestion
        - Penalize timing violations
        - Penalize thermal hotspots
        - Reward area utilization
        """
        reward = 0
        
        # 1. Wire Length (negative reward)
        wire_length = self._calculate_wire_length()
        self.total_wire_length = wire_length
        reward -= wire_length * 0.01  # Scale factor
        
        # 2. Congestion penalty
        congestion = self._calculate_congestion()
        self.max_congestion = congestion
        reward -= congestion * 0.5
        
        # 3. Timing violations
        timing_violations = self._calculate_timing_violations()
        self.timing_violations = timing_violations
        reward -= timing_violations * 2.0  # Higher penalty for timing
        
        # 4. Thermal hotspots
        thermal_score = self._calculate_thermal_score()
        self.thermal_hotspots = thermal_score
        reward -= thermal_score * 1.0
        
        # 5. Area utilization bonus
        utilization = len(self.placed_blocks) / self.num_blocks
        reward += utilization * 10
        
        # 6. Completion bonus
        if len(self.placed_blocks) == self.num_blocks:
            reward += 100
        
        return reward
    
    def _calculate_wire_length(self) -> float:
        """Calculate Half-Perimeter Wire Length (HPWL)"""
        total_length = 0
        
        for edge in self.netlist.edges(data=True):
            block1, block2, data = edge
            
            # Only calculate if both blocks placed
            if block1 in self.block_positions and block2 in self.block_positions:
                x1, y1 = self.block_positions[block1]
                x2, y2 = self.block_positions[block2]
                
                # HPWL = |x1 - x2| + |y1 - y2|
                hpwl = abs(x1 - x2) + abs(y1 - y2)
                weight = data['weight']
                total_length += hpwl * weight
        
        return total_length
    
    def _calculate_congestion(self) -> float:
        """Calculate routing congestion using grid-based model"""
        # Divide canvas into routing grid
        grid_size = 10
        routing_grid = np.zeros((
            self.canvas_width // grid_size,
            self.canvas_height // grid_size
        ))
        
        # For each net, add routing demand to grid cells
        for edge in self.netlist.edges():
            block1, block2 = edge
            
            if block1 in self.block_positions and block2 in self.block_positions:
                x1, y1 = self.block_positions[block1]
                x2, y2 = self.block_positions[block2]
                
                # Simple L-shaped routing
                # Horizontal segment
                for x in range(min(x1, x2), max(x1, x2)):
                    grid_x = x // grid_size
                    grid_y = y1 // grid_size
                    if grid_x < routing_grid.shape[0] and grid_y < routing_grid.shape[1]:
                        routing_grid[grid_x, grid_y] += 1
                
                # Vertical segment
                for y in range(min(y1, y2), max(y1, y2)):
                    grid_x = x2 // grid_size
                    grid_y = y // grid_size
                    if grid_x < routing_grid.shape[0] and grid_y < routing_grid.shape[1]:
                        routing_grid[grid_x, grid_y] += 1
        
        # Max congestion
        return np.max(routing_grid)
    
    def _calculate_timing_violations(self) -> int:
        """Count critical path timing violations"""
        violations = 0
        critical_threshold = 0.7  # Blocks with criticality > 0.7
        max_wire_length = 50  # Maximum wire length for critical paths
        
        for edge in self.netlist.edges(data=True):
            block1, block2, data = edge
            
            if block1 in self.block_positions and block2 in self.block_positions:
                # Check if path is timing-critical
                if (self.blocks[block1]['timing_criticality'] > critical_threshold or
                    self.blocks[block2]['timing_criticality'] > critical_threshold):
                    
                    x1, y1 = self.block_positions[block1]
                    x2, y2 = self.block_positions[block2]
                    wire_length = abs(x1 - x2) + abs(y1 - y2)
                    
                    if wire_length > max_wire_length:
                        violations += 1
        
        return violations
    
    def _calculate_thermal_score(self) -> float:
        """Calculate thermal hotspot penalty using power density map"""
        # Create power density grid
        power_grid = np.zeros((self.canvas_width, self.canvas_height))
        
        for block_id in self.placed_blocks:
            x, y = self.block_positions[block_id]
            block = self.blocks[block_id]
            
            # Add power density to grid
            for i in range(x, x + block['width']):
                for j in range(y, y + block['height']):
                    power_grid[i, j] = block['power_density']
        
        # Simulate thermal diffusion (simple averaging)
        thermal_grid = power_grid.copy()
        for _ in range(3):  # 3 iterations of diffusion
            thermal_grid = 0.25 * (
                np.roll(thermal_grid, 1, axis=0) +
                np.roll(thermal_grid, -1, axis=0) +
                np.roll(thermal_grid, 1, axis=1) +
                np.roll(thermal_grid, -1, axis=1)
            )
        
        # Count hotspots (temperature > threshold)
        hotspot_threshold = 1.5
        hotspots = np.sum(thermal_grid > hotspot_threshold)
        
        return hotspots
    
    def render(self, mode='human'):
        """Visualize current placement"""
        import matplotlib.pyplot as plt
        
        fig, ax = plt.subplots(figsize=(10, 10))
        
        # Draw canvas
        ax.set_xlim(0, self.canvas_width)
        ax.set_ylim(0, self.canvas_height)
        ax.set_aspect('equal')
        
        # Draw placed blocks
        for block_id in self.placed_blocks:
            x, y = self.block_positions[block_id]
            block = self.blocks[block_id]
            
            rect = plt.Rectangle(
                (x, y),
                block['width'],
                block['height'],
                fill=True,
                facecolor=plt.cm.viridis(block['power_density'] / 2.0),
                edgecolor='black',
                linewidth=2
            )
            ax.add_patch(rect)
            
            # Add label
            ax.text(
                x + block['width'] / 2,
                y + block['height'] / 2,
                f"B{block_id}",
                ha='center',
                va='center',
                fontsize=8,
                color='white',
                weight='bold'
            )
        
        # Draw nets
        for edge in self.netlist.edges():
            block1, block2 = edge
            
            if block1 in self.block_positions and block2 in self.block_positions:
                x1, y1 = self.block_positions[block1]
                x2, y2 = self.block_positions[block2]
                
                # Center of blocks
                b1 = self.blocks[block1]
                b2 = self.blocks[block2]
                cx1 = x1 + b1['width'] / 2
                cy1 = y1 + b1['height'] / 2
                cx2 = x2 + b2['width'] / 2
                cy2 = y2 + b2['height'] / 2
                
                ax.plot([cx1, cx2], [cy1, cy2], 'r-', alpha=0.3, linewidth=0.5)
        
        ax.set_title(f"Chip Floorplan - {len(self.placed_blocks)}/{self.num_blocks} blocks placed")
        ax.set_xlabel("X Position")
        ax.set_ylabel("Y Position")
        
        plt.grid(True, alpha=0.3)
        plt.show()

# ppo_agent.py
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.distributions import Categorical
import numpy as np
from typing import List, Tuple
 
class PolicyNetwork(nn.Module):
    """
    Actor network that outputs action probabilities
    """
    def __init__(self, state_dim: int, action_dim: int, hidden_dim: int = 256):
        super(PolicyNetwork, self).__init__()
        
        self.fc1 = nn.Linear(state_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.fc3 = nn.Linear(hidden_dim, action_dim)
        
        self.dropout = nn.Dropout(0.2)
        
    def forward(self, state: torch.Tensor) -> torch.Tensor:
        x = F.relu(self.fc1(state))
        x = self.dropout(x)
        x = F.relu(self.fc2(x))
        x = self.dropout(x)
        action_probs = F.softmax(self.fc3(x), dim=-1)
        return action_probs
 
class ValueNetwork(nn.Module):
    """
    Critic network that estimates state value
    """
    def __init__(self, state_dim: int, hidden_dim: int = 256):
        super(ValueNetwork, self).__init__()
        
        self.fc1 = nn.Linear(state_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.fc3 = nn.Linear(hidden_dim, 1)
        
        self.dropout = nn.Dropout(0.2)
        
    def forward(self, state: torch.Tensor) -> torch.Tensor:
        x = F.relu(self.fc1(state))
        x = self.dropout(x)
        x = F.relu(self.fc2(x))
        x = self.dropout(x)
        value = self.fc3(x)
        return value
 
class PPOAgent:
    """
    Proximal Policy Optimization agent for chip floorplanning
    """
    def __init__(
        self,
        state_dim: int,
        action_dim: int,
        lr: float = 3e-4,
        gamma: float = 0.99,
        epsilon: float = 0.2,
        c1: float = 1.0,  # Value loss coefficient
        c2: float = 0.01,  # Entropy bonus coefficient
        device: str = 'cuda' if torch.cuda.is_available() else 'cpu'
    ):
        self.device = device
        self.gamma = gamma
        self.epsilon = epsilon
        self.c1 = c1
        self.c2 = c2
        
        # Networks
        self.policy = PolicyNetwork(state_dim, action_dim).to(device)
        self.value = ValueNetwork(state_dim).to(device)
        
        # Optimizers
        self.policy_optimizer = optim.Adam(self.policy.parameters(), lr=lr)
        self.value_optimizer = optim.Adam(self.value.parameters(), lr=lr)
        
        # Experience buffer
        self.states = []
        self.actions = []
        self.rewards = []
        self.dones = []
        self.log_probs = []
        
    def select_action(self, state: np.ndarray, explore: bool = True) -> Tuple[int, float]:
        """Select action using policy network"""
        state_tensor = torch.FloatTensor(state).unsqueeze(0).to(self.device)
        
        with torch.no_grad():
            action_probs = self.policy(state_tensor)
        
        if explore:
            # Sample from distribution
            dist = Categorical(action_probs)
            action = dist.sample()
            log_prob = dist.log_prob(action)
        else:
            # Greedy action
            action = torch.argmax(action_probs, dim=1)
            log_prob = torch.log(action_probs.squeeze(0)[action])
        
        return action.item(), log_prob.item()
    
    def store_transition(
        self,
        state: np.ndarray,
        action: int,
        reward: float,
        done: bool,
        log_prob: float
    ):
        """Store experience in buffer"""
        self.states.append(state)
        self.actions.append(action)
        self.rewards.append(reward)
        self.dones.append(done)
        self.log_probs.append(log_prob)
    
    def compute_returns(self) -> List[float]:
        """Compute discounted returns (Monte Carlo)"""
        returns = []
        R = 0
        
        for reward, done in zip(reversed(self.rewards), reversed(self.dones)):
            if done:
                R = 0
            R = reward + self.gamma * R
            returns.insert(0, R)
        
        return returns
    
    def compute_advantages(self, returns: List[float]) -> np.ndarray:
        """Compute advantage estimates using GAE"""
        states_tensor = torch.FloatTensor(self.states).to(self.device)
        
        with torch.no_grad():
            values = self.value(states_tensor).squeeze().cpu().numpy()
        
        returns = np.array(returns)
        advantages = returns - values
        
        # Normalize advantages
        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
        
        return advantages
    
    def update(self, epochs: int = 10, batch_size: int = 64):
        """Update policy and value networks using PPO"""
        # Compute returns and advantages
        returns = self.compute_returns()
        advantages = self.compute_advantages(returns)
        
        # Convert to tensors
        states_tensor = torch.FloatTensor(self.states).to(self.device)
        actions_tensor = torch.LongTensor(self.actions).to(self.device)
        old_log_probs = torch.FloatTensor(self.log_probs).to(self.device)
        returns_tensor = torch.FloatTensor(returns).to(self.device)
        advantages_tensor = torch.FloatTensor(advantages).to(self.device)
        
        # PPO update for multiple epochs
        for _ in range(epochs):
            # Shuffle data
            indices = np.random.permutation(len(self.states))
            
            for start in range(0, len(self.states), batch_size):
                end = start + batch_size
                batch_indices = indices[start:end]
                
                batch_states = states_tensor[batch_indices]
                batch_actions = actions_tensor[batch_indices]
                batch_old_log_probs = old_log_probs[batch_indices]
                batch_returns = returns_tensor[batch_indices]
                batch_advantages = advantages_tensor[batch_indices]
                
                # Compute new action probabilities
                action_probs = self.policy(batch_states)
                dist = Categorical(action_probs)
                new_log_probs = dist.log_prob(batch_actions)
                entropy = dist.entropy().mean()
                
                # Compute ratio for PPO
                ratio = torch.exp(new_log_probs - batch_old_log_probs)
                
                # Clipped surrogate objective
                surr1 = ratio * batch_advantages
                surr2 = torch.clamp(ratio, 1 - self.epsilon, 1 + self.epsilon) * batch_advantages
                policy_loss = -torch.min(surr1, surr2).mean()
                
                # Add entropy bonus for exploration
                policy_loss -= self.c2 * entropy
                
                # Update policy
                self.policy_optimizer.zero_grad()
                policy_loss.backward()
                torch.nn.utils.clip_grad_norm_(self.policy.parameters(), 0.5)
                self.policy_optimizer.step()
                
                # Value loss
                values = self.value(batch_states).squeeze()
                value_loss = F.mse_loss(values, batch_returns)
                
                # Update value network
                self.value_optimizer.zero_grad()
                value_loss.backward()
                torch.nn.utils.clip_grad_norm_(self.value.parameters(), 0.5)
                self.value_optimizer.step()
        
        # Clear buffer
        self.states = []
        self.actions = []
        self.rewards = []
        self.dones = []
        self.log_probs = []
        
        return policy_loss.item(), value_loss.item()
    
    def save(self, path: str):
        """Save model weights"""
        torch.save({
            'policy_state_dict': self.policy.state_dict(),
            'value_state_dict': self.value.state_dict(),
            'policy_optimizer_state_dict': self.policy_optimizer.state_dict(),
            'value_optimizer_state_dict': self.value_optimizer.state_dict(),
        }, path)
    
    def load(self, path: str):
        """Load model weights"""
        checkpoint = torch.load(path, map_location=self.device)
        self.policy.load_state_dict(checkpoint['policy_state_dict'])
        self.value.load_state_dict(checkpoint['value_state_dict'])
        self.policy_optimizer.load_state_dict(checkpoint['policy_optimizer_state_dict'])
        self.value_optimizer.load_state_dict(checkpoint['value_optimizer_state_dict'])

# train.py
import numpy as np
import matplotlib.pyplot as plt
from chip_env import ChipFloorplanEnv
from ppo_agent import PPOAgent
import wandb  # For experiment tracking
 
def train_agent(
    env: ChipFloorplanEnv,
    agent: PPOAgent,
    num_episodes: int = 1000,
    max_steps: int = 100,
    update_interval: int = 10,
    save_interval: int = 100
):
    """Train PPO agent on chip floorplanning"""
    
    # Initialize W&B for experiment tracking
    wandb.init(
        project="chip-floorplanning",
        config={
            "num_episodes": num_episodes,
            "max_steps": max_steps,
            "learning_rate": 3e-4,
            "gamma": 0.99,
            "epsilon": 0.2
        }
    )
    
    episode_rewards = []
    episode_wire_lengths = []
    episode_congestions = []
    
    for episode in range(num_episodes):
        state = env.reset()
        episode_reward = 0
        
        for step in range(max_steps):
            # Select action
            action, log_prob = agent.select_action(state, explore=True)
            
            # Decompose action (assuming flattened action space)
            # For simplicity, map action to (block_id, x, y)
            block_id = action % env.num_blocks
            position = action // env.num_blocks
            x = position % env.canvas_width
            y = position // env.canvas_width
            
            # Take step in environment
            next_state, reward, done, info = env.step([block_id, x, y])
            
            # Store transition
            agent.store_transition(state, action, reward, done, log_prob)
            
            episode_reward += reward
            state = next_state
            
            if done:
                break
        
        # Update agent every N episodes
        if (episode + 1) % update_interval == 0:
            policy_loss, value_loss = agent.update()
            
            wandb.log({
                "policy_loss": policy_loss,
                "value_loss": value_loss,
                "episode": episode
            })
        
        # Log metrics
        episode_rewards.append(episode_reward)
        episode_wire_lengths.append(info.get('wire_length', 0))
        episode_congestions.append(info.get('congestion', 0))
        
        wandb.log({
            "episode_reward": episode_reward,
            "wire_length": info.get('wire_length', 0),
            "congestion": info.get('congestion', 0),
            "timing_violations": info.get('timing_violations', 0),
            "thermal_hotspots": info.get('thermal_hotspots', 0),
            "blocks_placed": info.get('blocks_placed', 0),
            "episode": episode
        })
        
        # Print progress
        if (episode + 1) % 10 == 0:
            avg_reward = np.mean(episode_rewards[-10:])
            avg_wire_length = np.mean(episode_wire_lengths[-10:])
            print(f"Episode {episode + 1}/{num_episodes}")
            print(f"  Avg Reward (last 10): {avg_reward:.2f}")
            print(f"  Avg Wire Length: {avg_wire_length:.2f}")
            print(f"  Blocks Placed: {info.get('blocks_placed', 0)}/{env.num_blocks}")
        
        # Save checkpoint
        if (episode + 1) % save_interval == 0:
            agent.save(f"checkpoints/ppo_agent_ep{episode+1}.pth")
            print(f"  → Saved checkpoint at episode {episode + 1}")
    
    # Plot training curves
    plot_training_curves(episode_rewards, episode_wire_lengths, episode_congestions)
    
    wandb.finish()
    
    return agent
 
def plot_training_curves(rewards, wire_lengths, congestions):
    """Plot training metrics"""
    fig, axes = plt.subplots(3, 1, figsize=(12, 10))
    
    # Reward curve
    axes[0].plot(rewards, alpha=0.3, color='blue')
    axes[0].plot(np.convolve(rewards, np.ones(50)/50, mode='valid'), color='blue', linewidth=2)
    axes[0].set_xlabel('Episode')
    axes[0].set_ylabel('Total Reward')
    axes[0].set_title('Training Reward')
    axes[0].grid(True, alpha=0.3)
    
    # Wire length curve
    axes[1].plot(wire_lengths, alpha=0.3, color='red')
    axes[1].plot(np.convolve(wire_lengths, np.ones(50)/50, mode='valid'), color='red', linewidth=2)
    axes[1].set_xlabel('Episode')
    axes[1].set_ylabel('Total Wire Length')
    axes[1].set_title('Wire Length Optimization')
    axes[1].grid(True, alpha=0.3)
    
    # Congestion curve
    axes[2].plot(congestions, alpha=0.3, color='green')
    axes[2].plot(np.convolve(congestions, np.ones(50)/50, mode='valid'), color='green', linewidth=2)
    axes[2].set_xlabel('Episode')
    axes[2].set_ylabel('Max Congestion')
    axes[2].set_title('Congestion Minimization')
    axes[2].grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('training_curves.png', dpi=300)
    plt.show()
 
if __name__ == "__main__":
    # Create environment
    env = ChipFloorplanEnv(
        canvas_width=100,
        canvas_height=100,
        num_blocks=20,
        netlist_path="data/netlist.txt"
    )
    
    # Create agent
    state_dim = env.observation_space.shape[0]
    action_dim = env.num_blocks * env.canvas_width * env.canvas_height
    
    agent = PPOAgent(
        state_dim=state_dim,
        action_dim=action_dim,
        lr=3e-4,
        gamma=0.99,
        epsilon=0.2
    )
    
    # Train
    trained_agent = train_agent(
        env=env,
        agent=agent,
        num_episodes=1000,
        max_steps=100,
        update_interval=10,
        save_interval=100
    )
    
    # Evaluate
    print("\n=== Evaluation ===")
    state = env.reset()
    done = False
    
    while not done:
        action, _ = agent.select_action(state, explore=False)
        block_id = action % env.num_blocks
        position = action // env.num_blocks
        x = position % env.canvas_width
        y = position // env.canvas_width
        
        state, reward, done, info = env.step([block_id, x, y])
    
    # Visualize final placement
    env.render()
    
    print(f"\nFinal Metrics:")
    print(f"  Wire Length: {info['wire_length']:.2f}")
    print(f"  Congestion: {info['congestion']:.2f}")
    print(f"  Timing Violations: {info['timing_violations']}")
    print(f"  Thermal Hotspots: {info['thermal_hotspots']}")

# Incremental reward for each valid placement
reward += 1.0
 
# Bonus for reducing wire length compared to previous step
if wire_length < previous_wire_length:
    reward += 5.0
 
# Penalty for increasing congestion
if congestion > previous_congestion:
    reward -= 2.0

def get_valid_actions(self, state):
    """Return mask of valid actions"""
    valid_mask = np.zeros(self.action_space.n)
    
    for action in range(self.action_space.n):
        block_id, x, y = self.decode_action(action)
        
        if block_id not in self.placed_blocks and self._is_valid_placement(block_id, x, y):
            valid_mask[action] = 1
    
    return valid_mask
 
# In agent's select_action()
valid_mask = env.get_valid_actions(state)
action_probs = action_probs * valid_mask
action_probs = action_probs / action_probs.sum()  # Renormalize

# Start with small designs (5 blocks)
# Gradually increase to full design (20 blocks)
def get_curriculum_stage(episode):
    if episode < 200:
        return 5  # 5 blocks
    elif episode < 500:
        return 10  # 10 blocks
    elif episode < 800:
        return 15  # 15 blocks
    else:
        return 20  # Full design

class HierarchicalAgent:
    def __init__(self):
        self.high_level_policy = ClusteringPolicy()  # Groups blocks
        self.low_level_policy = PlacementPolicy()    # Places blocks
    
    def select_action(self, state):
        # First, decide which cluster to place
        cluster = self.high_level_policy.select_cluster(state)
        
        # Then, decide where to place blocks in cluster
        placement = self.low_level_policy.select_placement(state, cluster)
        
        return placement

import torch_geometric.nn as pyg_nn
 
class GNNPolicy(nn.Module):
    def __init__(self, node_dim, edge_dim, hidden_dim):
        super().__init__()
        self.conv1 = pyg_nn.GCNConv(node_dim, hidden_dim)
        self.conv2 = pyg_nn.GCNConv(hidden_dim, hidden_dim)
        self.fc = nn.Linear(hidden_dim, 2)  # Output (x, y)
    
    def forward(self, x, edge_index, edge_attr):
        x = F.relu(self.conv1(x, edge_index))
        x = F.relu(self.conv2(x, edge_index))
        placement = self.fc(x)
        return placement

# Load pretrained weights
agent.load("pretrained_openroad.pth")
 
# Fine-tune on AMD design
for episode in range(100):  # Only 100 episodes needed
    # Train on AMD netlist
    state = amd_env.reset()
    ...