hybrid_policy_training.py

import os
import gymnasium as gym
import numpy as np
import torch
import torch.nn as nn
from stable_baselines3 import PPO
from stable_baselines3.common.callbacks import BaseCallback, CheckpointCallback
import wandb

# Initialize wandb with API key
wandb.login(key="bd346658e170bea330a512ff4f7348be3d9e5547")

class HybridPolicy(nn.Module):
    def __init__(self, state_dim, action_dim):
        super().__init__()
        self.shared_network = nn.Sequential(
            nn.Linear(state_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU()
        )
        
        # Policy-specific heads
        self.walking_head = nn.Sequential(
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, action_dim)
        )
        
        self.running_head = nn.Sequential(
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, action_dim)
        )
        
        # Gating network to decide policy mixing weights
        self.gate_network = nn.Sequential(
            nn.Linear(state_dim, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Sigmoid()
        )
        
    def forward(self, state):
        shared_features = self.shared_network(state)
        
        # Get actions from both policies
        walking_action = self.walking_head(shared_features)
        running_action = self.running_head(shared_features)
        
        # Calculate mixing weight
        gate_value = self.gate_network(state)
        
        # Blend actions
        blended_action = gate_value * walking_action + (1 - gate_value) * running_action
        return blended_action

class WandbCallback(BaseCallback):
    def __init__(self, verbose=0):
        super().__init__(verbose)
        self.episode_rewards = []
        self.episode_lengths = []
        
    def _on_step(self):
        # Log episode rewards
        if self.locals.get("dones")[0]:
            episode_reward = self.locals.get("rewards")[0]
            episode_length = self.locals.get("n_steps")
            wandb.log({
                "episode_reward": episode_reward,
                "episode_length": episode_length,
                "timesteps": self.num_timesteps
            })
        return True

def create_env(task="walk"):
    """Create environment with specific task parameters."""
    env = gym.make("Humanoid-v4")
    if task == "run":
        env.unwrapped.target_velocity = 5.0  # Higher velocity for running
    else:
        env.unwrapped.target_velocity = 2.0  # Lower velocity for walking
    return env

def train_individual_policy(task="walk", total_timesteps=200000):
    """Train a single policy for either walking or running."""
    env = create_env(task)
    
    model = PPO(
        "MlpPolicy",
        env,
        verbose=1,
        learning_rate=3e-4,
        n_steps=2048,
        batch_size=64,
        n_epochs=10,
        gamma=0.99,
        device="cuda" if torch.cuda.is_available() else "cpu"
    )
    
    wandb.init(
        project="hybrid-policy-training",
        name=f"{task}-policy-training",
        config={
            "task": task,
            "total_timesteps": total_timesteps,
            "algorithm": "PPO"
        }
    )
    
    # Create checkpoint directory if it doesn't exist
    os.makedirs("./checkpoints", exist_ok=True)
    
    callback = WandbCallback()
    checkpoint_callback = CheckpointCallback(
        save_freq=50000,
        save_path="./checkpoints/",
        name_prefix=f"{task}_policy"
    )
    callbacks = [callback, checkpoint_callback]
    
    model.learn(total_timesteps=total_timesteps, callback=callbacks)
    model.save(f"{task}_policy_final")
    wandb.finish()
    return model

def train_hybrid_policy(walking_policy, running_policy, total_timesteps=100000):
    """Train the hybrid policy using both walking and running policies."""
    env = create_env("hybrid")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # Initialize hybrid policy
    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.shape[0]
    hybrid_policy = HybridPolicy(state_dim, action_dim).to(device)
    
    optimizer = torch.optim.Adam(hybrid_policy.parameters(), lr=3e-4)
    
    wandb.init(
        project="hybrid-policy-training",
        name="hybrid-policy-training",
        config={
            "total_timesteps": total_timesteps
        }
    )
    
    # Training loop
    for step in range(total_timesteps):
        state = env.reset()[0]
        done = False
        episode_reward = 0
        
        while not done:
            state_tensor = torch.FloatTensor(state).unsqueeze(0).to(device)
            action = hybrid_policy(state_tensor)
            action_np = action.cpu().detach().numpy().squeeze()
            
            next_state, reward, done, truncated, info = env.step(action_np)
            episode_reward += reward
            
            # Compute loss using reward as a tensor
            reward_tensor = torch.tensor(reward, dtype=torch.float32, device=device)
            loss = -torch.mean(action * reward_tensor)
            
            # Update policy
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            state = next_state
            
            if done:
                wandb.log({
                    "episode_reward": episode_reward,
                    "loss": loss.item(),
                    "step": step
                })
                break
    
    wandb.finish()
    return hybrid_policy

def main():
    # Train individual policies
    print("Training walking policy...")
    walking_policy = train_individual_policy("walk")
    
    print("Training running policy...")
    running_policy = train_individual_policy("run")
    
    print("Training hybrid policy...")
    hybrid_policy = train_hybrid_policy(walking_policy, running_policy)
    
    # Save the final hybrid policy
    torch.save(hybrid_policy.state_dict(), "hybrid_policy.pth")

if __name__ == "__main__":
    main()