Update app.py

app.py

@@ -5,49 +5,42 @@ import torch.optim as optim
 from torch.distributions import Categorical
 import gradio as gr
 import matplotlib.pyplot as plt
-from collections import defaultdict
-import random
+import io
+from PIL import Image
 
 # ==========================================
-# 1. БЕСКОНЕЧНЫЙ ПРОЦЕДУРНЫЙ МИР
+# 1. БЕСКОНЕЧНЫЙ МИР
 # ==========================================
 class InfiniteWorld:
-    """
-    Мир хранится в словаре чанков. 
-    Координаты не ограничены. 
-    Биомы генерируются детерминировано по hash координат.
-    """
     CHUNK_SIZE = 16
-    VIEW_RADIUS = 8  # Агент видит 16x16 вокруг себя
+    VIEW_RADIUS = 8
     
     def __init__(self, seed=42):
         self.seed = seed
-        self.chunks = {}  # (cx, cy) -> np.array(CHUNK_SIZE, CHUNK_SIZE)
+        self.chunks = {}
         self.agent_pos = [0, 0]
         self.steps = 0
         self.max_steps = 1000
         
-    def _get_chunk(self, cx: int, cy: int) -> np.ndarray:
+    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  # "Природные" блоки
+            chunk[noise > 0.7] = 1.0
             self.chunks[(cx, cy)] = chunk
         return self.chunks[(cx, cy)]
     
-    def _world_coords(self, x: int, y: int):
+    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: int, y: int) -> float:
+    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: int, y: int, val: float):
+    def set_block(self, x, y, val):
         cx, cy, lx, ly = self._world_coords(x, y)
         self._get_chunk(cx, cy)[lx, ly] = val
     
@@ -57,10 +50,8 @@ class InfiniteWorld:
         return self._get_obs()
     
     def _get_obs(self):
-        """Возвращает локальный патч 16x16x3 вокруг агента"""
         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
@@ -68,52 +59,40 @@ class InfiniteWorld:
                 px = dx + self.VIEW_RADIUS
                 py = dy + self.VIEW_RADIUS
                 patch[px, py, 0] = block
-                # Канал 1: расстояние до центра (позиционный энкодинг)
                 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: int):
+    def step(self, action):
         self.steps += 1
         reward = -0.005
         done = self.steps >= self.max_steps
         
-        ACTIONS = {
-            0: (-1, 0), 1: (1, 0), 2: (0, -1), 3: (0, 1),  # Move N/S/W/E
-            4: (0, 0),   # Build
-            5: (0, 0),   # Dig
-        }
-        
-        if action < 4:
-            dx, dy = ACTIONS[action]
-            self.agent_pos[0] += dx
-            self.agent_pos[1] += dy
-        elif action == 4:  # Build
+        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:  # Dig
+        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, {"pos": tuple(self.agent_pos)}
+        return self._get_obs(), reward, done, {}
 
 
 # ==========================================
-# 2. PPO AGENT С ПАМЯТЬЮ
+# 2. PPO AGENT
 # ==========================================
 class PPOAgent(nn.Module):
-    def __init__(self, action_space=6, hidden_dim=256):
+    def __init__(self):
         super().__init__()
-        self.action_space = action_space
-        
-        # Vision encoder (обрабатывает локальный патч 16x16)
         self.encoder = nn.Sequential(
             nn.Conv2d(3, 32, 3, padding=1), nn.ReLU(),
             nn.Conv2d(32, 64, 3, stride=2, padding=1), nn.ReLU(),
@@ -121,21 +100,14 @@ class PPOAgent(nn.Module):
             nn.AdaptiveAvgPool2d((4, 4)),
             nn.Flatten()
         )
-        
-        # GRU для памяти (агент помнит, что уже видел/строил)
-        self.gru = nn.GRUCell(64 * 4 * 4, hidden_dim)
-        self.hidden = None
-        
-        # Actor-Critic heads
-        self.actor = nn.Linear(hidden_dim, action_space)
-        self.critic = nn.Linear(hidden_dim, 1)
+        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)
-        logits = self.actor(h)
-        value = self.critic(h)
-        return logits, value, h
+        return self.actor(h), self.critic(h), h
     
     def act(self, obs, hidden=None):
         with torch.no_grad():
@@ -146,52 +118,42 @@ class PPOAgent(nn.Module):
 
 
 # ==========================================
-# 3. PPO TRAINING LOOP
+# 3. ОБУЧЕНИЕ
 # ==========================================
-def train_ppo(episodes=300, steps_per_update=256, lr=3e-4):
+def train_ppo(episodes=100):
     env = InfiniteWorld()
     agent = PPOAgent()
-    optimizer = optim.Adam(agent.parameters(), lr=lr)
-    
-    reward_history = []
+    optimizer = optim.Adam(agent.parameters(), lr=3e-4)
     
     for ep in range(episodes):
         obs = env.reset()
         hidden = None
-        episode_rewards = []
+        buffers = {'obs': [], 'actions': [], 'log_probs': [], 'rewards': [], 'values': []}
         
-        # Rollout buffer
-        buffers = {'obs': [], 'actions': [], 'log_probs': [], 'rewards': [], 'values': [], 'hiddens': []}
-        
-        for _ in range(steps_per_update):
+        for _ in range(256):
             obs_t = torch.FloatTensor(obs)
             action, log_prob, value, hidden = agent.act(obs_t, hidden)
-            next_obs, reward, done, info = env.step(action)
+            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)
-            buffers['hiddens'].append(hidden)
-            episode_rewards.append(reward)
             
             obs = next_obs
             if done:
                 obs = env.reset()
                 hidden = None
         
-        # GAE computation
-        returns = []
-        advantages = []
-        R = 0
-        A = 0
-        gamma, lam = 0.99, 0.95
-        
+        # GAE
+        returns, advantages = [], []
+        R, A = 0, 0
         for i in reversed(range(len(buffers['rewards']))):
-            R = buffers['rewards'][i] + gamma * R
-            delta = buffers['rewards'][i] + gamma * (buffers['values'][i+1].item() if i < len(buffers['values'])-1 else 0) - buffers['values'][i].item()
-            A = delta + gamma * lam * A
+            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)
         
@@ -199,7 +161,6 @@ def train_ppo(episodes=300, steps_per_update=256, lr=3e-4):
         advantages = torch.FloatTensor(advantages)
         advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
         
-        # PPO Update (4 epochs)
         obs_batch = torch.stack(buffers['obs'])
         actions_batch = torch.LongTensor(buffers['actions'])
         old_log_probs = torch.stack(buffers['log_probs']).detach()
@@ -208,112 +169,73 @@ def train_ppo(episodes=300, steps_per_update=256, lr=3e-4):
             logits, values, _ = agent.forward(obs_batch)
             dist = Categorical(logits=logits)
             new_log_probs = dist.log_prob(actions_batch)
-            entropy = dist.entropy().mean()
-            
             ratio = (new_log_probs - old_log_probs).exp()
-            surr1 = ratio * advantages
-            surr2 = torch.clamp(ratio, 0.8, 1.2) * advantages
-            
-            actor_loss = -torch.min(surr1, surr2).mean()
-            critic_loss = (returns - values.squeeze()).pow(2).mean()
-            loss = actor_loss + 0.5 * critic_loss - 0.01 * entropy
+            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()
         
-        avg_reward = np.mean(episode_rewards)
-        reward_history.append(avg_reward)
         if ep % 20 == 0:
-            print(f"Ep {ep} | Avg Reward: {avg_reward:.3f} | Chunks explored: {len(env.chunks)}")
+            print(f"Ep {ep} | Chunks: {len(env.chunks)}")
     
-    return agent, reward_history
+    return agent
 
 
 # ==========================================
-# 4. HF SPACE DEMO
-# ==========================================
-# ==========================================
-# 4. HF SPACE DEMO (ИСПРАВЛЕНО ДЛЯ HF SPACES)
+# 4. ГРАФИЧЕСКИЙ ИНТЕРФЕЙС (ИСПРАВЛЕНО)
 # ==========================================
-def create_demo():
-    print("🏗️ Обучение PPO-агента в бесконечном мире...")
-    agent, history = train_ppo(episodes=100)
-    agent.eval()
-    
-    def run_exploration(n_steps: int = 200):
-        """Функция с явной типизацией для избежания бага schema в Gradio"""
-        env = InfiniteWorld(seed=random.randint(0, 99999))
-        obs = env.reset()
-        hidden = None
-        frames = []
-        positions = []
-        
-        with torch.no_grad():
-            for _ in range(min(n_steps, 500)):  # Ограничиваем макс шаги
-                fig, ax = plt.subplots(figsize=(5, 5))
-                ax.imshow(obs)
-                ax.set_title(f"Pos: {env.agent_pos} | Chunks: {len(env.chunks)}")
-                ax.axis('off')
-                frames.append(fig)
-                plt.close(fig)
-                positions.append(tuple(env.agent_pos))
-                
-                obs_t = torch.FloatTensor(obs)
-                action, _, _, hidden = agent.act(obs_t, hidden)
-                obs, _, done, _ = env.step(action)
-                if done:
-                    break
-        
-        # Карта траектории
-        if positions:
-            fig2, ax2 = plt.subplots(figsize=(6, 6))
-            xs, ys = zip(*positions)
-            ax2.scatter(xs, ys, c=range(len(positions)), cmap='viridis', s=1)
-            ax2.set_title("Траектория исследования")
-            ax2.set_aspect('equal')
-            frames.append(fig2)
-            plt.close(fig2)
-        
-        return frames
+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
     
-    # Явное указание типов компонентов устраняет баг schema
-    with gr.Blocks(title="Infinite Builder Agent") as demo:
-        gr.Markdown("# 🌍 Бесконечный мир: PPO-агент с нуля")
-        gr.Markdown("Агент видит только локальный патч 16×16, имеет GRU-память и исследует процедурный мир.")
-        
-        steps_input = gr.Slider(
-            minimum=50, 
-            maximum=500, 
-            step=50, 
-            value=200, 
-            label="Количество шагов исследования"
-        )
-        btn = gr.Button("▶️ Запустить исследование", variant="primary")
-        gallery = gr.Gallery(
-            label="Процесс строительства и исследования", 
-            columns=4, 
-            height="auto",
-            object_fit="contain"
-        )
-        
-        btn.click(
-            fn=run_exploration, 
-            inputs=[steps_input],
-            outputs=[gallery]
-        )
+    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 demo
+    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 = create_demo()
-    # КРИТИЧНО: Для HF Spaces обязательно server_name="0.0.0.0"
-    # share=False на Spaces, т.к. прокси сам маршрутизирует трафик
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=True,
-        show_api=False
-    )
+    demo.launch(server_name="0.0.0.0", server_port=7860)