app.py

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions import Categorical
import gradio as gr
import matplotlib.pyplot as plt
import io
from PIL import Image

# ==========================================
# 1. БЕСКОНЕЧНЫЙ МИР
# ==========================================
class InfiniteWorld:
    CHUNK_SIZE = 16
    VIEW_RADIUS = 8
    
    def __init__(self, seed=42):
        self.seed = seed
        self.chunks = {}
        self.agent_pos = [0, 0]
        self.steps = 0
        self.max_steps = 1000
        
    def _get_chunk(self, cx, cy):
        if (cx, cy) not in self.chunks:
            rng = np.random.RandomState(hash((cx, cy, self.seed)) % (2**31))
            chunk = np.zeros((self.CHUNK_SIZE, self.CHUNK_SIZE), dtype=np.float32)
            noise = rng.rand(self.CHUNK_SIZE, self.CHUNK_SIZE)
            chunk[noise > 0.7] = 1.0
            self.chunks[(cx, cy)] = chunk
        return self.chunks[(cx, cy)]
    
    def _world_coords(self, x, y):
        cx, lx = divmod(x, self.CHUNK_SIZE)
        cy, ly = divmod(y, self.CHUNK_SIZE)
        return cx, cy, lx, ly
    
    def get_block(self, x, y):
        cx, cy, lx, ly = self._world_coords(x, y)
        return self._get_chunk(cx, cy)[lx, ly]
    
    def set_block(self, x, y, val):
        cx, cy, lx, ly = self._world_coords(x, y)
        self._get_chunk(cx, cy)[lx, ly] = val
    
    def reset(self):
        self.agent_pos = [0, 0]
        self.steps = 0
        return self._get_obs()
    
    def _get_obs(self):
        x, y = self.agent_pos
        patch = np.zeros((self.VIEW_RADIUS*2, self.VIEW_RADIUS*2, 3), dtype=np.float32)
        for dx in range(-self.VIEW_RADIUS, self.VIEW_RADIUS):
            for dy in range(-self.VIEW_RADIUS, self.VIEW_RADIUS):
                wx, wy = x + dx, y + dy
                block = self.get_block(wx, wy)
                px = dx + self.VIEW_RADIUS
                py = dy + self.VIEW_RADIUS
                patch[px, py, 0] = block
                patch[px, py, 1] = max(0, 1.0 - abs(dx)/self.VIEW_RADIUS)
                patch[px, py, 2] = max(0, 1.0 - abs(dy)/self.VIEW_RADIUS)
        patch[self.VIEW_RADIUS, self.VIEW_RADIUS, 1] = 1.0
        return patch
    
    def step(self, action):
        self.steps += 1
        reward = -0.005
        done = self.steps >= self.max_steps
        
        if action == 0: self.agent_pos[0] -= 1
        elif action == 1: self.agent_pos[0] += 1
        elif action == 2: self.agent_pos[1] -= 1
        elif action == 3: self.agent_pos[1] += 1
        elif action == 4:
            x, y = self.agent_pos
            if self.get_block(x, y) == 0:
                self.set_block(x, y, 1.0)
                reward = 1.0
        elif action == 5:
            x, y = self.agent_pos
            if self.get_block(x, y) == 1.0:
                self.set_block(x, y, 0.0)
                reward = 0.3
        
        return self._get_obs(), reward, done, {}


# ==========================================
# 2. PPO AGENT
# ==========================================
class PPOAgent(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(3, 32, 3, padding=1), nn.ReLU(),
            nn.Conv2d(32, 64, 3, stride=2, padding=1), nn.ReLU(),
            nn.Conv2d(64, 64, 3, stride=2, padding=1), nn.ReLU(),
            nn.AdaptiveAvgPool2d((4, 4)),
            nn.Flatten()
        )
        self.gru = nn.GRUCell(64 * 4 * 4, 256)
        self.actor = nn.Linear(256, 6)
        self.critic = nn.Linear(256, 1)
    
    def forward(self, obs, hidden=None):
        features = self.encoder(obs.permute(0, 3, 1, 2))
        h = self.gru(features, hidden)
        return self.actor(h), self.critic(h), h
    
    def act(self, obs, hidden=None):
        with torch.no_grad():
            logits, value, new_hidden = self.forward(obs.unsqueeze(0), hidden)
            dist = Categorical(logits=logits)
            action = dist.sample()
            return action.item(), dist.log_prob(action), value.squeeze(), new_hidden


# ==========================================
# 3. ОБУЧЕНИЕ
# ==========================================
def train_ppo(episodes=100):
    env = InfiniteWorld()
    agent = PPOAgent()
    optimizer = optim.Adam(agent.parameters(), lr=3e-4)
    
    for ep in range(episodes):
        obs = env.reset()
        hidden = None
        buffers = {'obs': [], 'actions': [], 'log_probs': [], 'rewards': [], 'values': []}
        
        for _ in range(256):
            obs_t = torch.FloatTensor(obs)
            action, log_prob, value, hidden = agent.act(obs_t, hidden)
            next_obs, reward, done, _ = env.step(action)
            
            buffers['obs'].append(obs_t)
            buffers['actions'].append(action)
            buffers['log_probs'].append(log_prob)
            buffers['rewards'].append(reward)
            buffers['values'].append(value)
            
            obs = next_obs
            if done:
                obs = env.reset()
                hidden = None
        
        # GAE
        returns, advantages = [], []
        R, A = 0, 0
        for i in reversed(range(len(buffers['rewards']))):
            R = buffers['rewards'][i] + 0.99 * R
            next_val = buffers['values'][i+1].item() if i < len(buffers['values'])-1 else 0
            delta = buffers['rewards'][i] + 0.99 * next_val - buffers['values'][i].item()
            A = delta + 0.99 * 0.95 * A
            returns.insert(0, R)
            advantages.insert(0, A)
        
        returns = torch.FloatTensor(returns)
        advantages = torch.FloatTensor(advantages)
        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
        
        obs_batch = torch.stack(buffers['obs'])
        actions_batch = torch.LongTensor(buffers['actions'])
        old_log_probs = torch.stack(buffers['log_probs']).detach()
        
        for _ in range(4):
            logits, values, _ = agent.forward(obs_batch)
            dist = Categorical(logits=logits)
            new_log_probs = dist.log_prob(actions_batch)
            ratio = (new_log_probs - old_log_probs).exp()
            surr = torch.min(ratio * advantages, torch.clamp(ratio, 0.8, 1.2) * advantages)
            loss = -surr.mean() + 0.5 * (returns - values.squeeze()).pow(2).mean() - 0.01 * dist.entropy().mean()
            
            optimizer.zero_grad()
            loss.backward()
            nn.utils.clip_grad_norm_(agent.parameters(), 0.5)
            optimizer.step()
        
        if ep % 20 == 0:
            print(f"Ep {ep} | Chunks: {len(env.chunks)}")
    
    return agent


# ==========================================
# 4. ГРАФИЧЕСКИЙ ИНТЕРФЕЙС (ИСПРАВЛЕНО)
# ==========================================
def fig_to_pil(fig):
    """Конвертирует matplotlib figure в PIL Image без schema-багов"""
    buf = io.BytesIO()
    fig.savefig(buf, format='png', bbox_inches='tight')
    buf.seek(0)
    img = Image.open(buf)
    plt.close(fig)
    return img


def run_simulation(n_steps):
    n_steps = int(n_steps)
    agent = run_simulation.agent
    env = InfiniteWorld(seed=np.random.randint(0, 99999))
    obs = env.reset()
    hidden = None
    
    images = []
    with torch.no_grad():
        for _ in range(min(n_steps, 300)):
            fig, ax = plt.subplots(figsize=(4, 4))
            ax.imshow(obs)
            ax.set_title(f"Pos: {env.agent_pos}")
            ax.axis('off')
            images.append(fig_to_pil(fig))
            
            obs_t = torch.FloatTensor(obs)
            action, _, _, hidden = agent.act(obs_t, hidden)
            obs, _, done, _ = env.step(action)
            if done:
                break
    
    return images


# Предобучаем модель один раз при загрузке
print("🏗️ Обучение агента...")
run_simulation.agent = train_ppo(episodes=80)
run_simulation.agent.eval()
print("✅ Обучение завершено!")


# Интерфейс БЕЗ типизации возврата, БЕЗ Gallery
with gr.Blocks(title="Infinite Builder") as demo:
    gr.Markdown("# 🌍 Бесконечный мир: PPO-агент")
    slider = gr.Slider(50, 300, value=100, step=50, label="Шагов")
    btn = gr.Button("▶️ Запустить")
    output = gr.Gallery(label="Результат", columns=4)
    btn.click(fn=run_simulation, inputs=[slider], outputs=[output])

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)