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)