""" Series 22, Post 4: World Models for Robots Animation: Model-free trial&error → RSSM latent transitions → Imagination rollouts → Real robot planning ~25 seconds at 24fps = 600 frames """ import numpy as np import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt from matplotlib.animation import FuncAnimation, FFMpegWriter from matplotlib.patches import FancyBboxPatch, FancyArrowPatch import matplotlib.patheffects as pe BG = '#0d1117' FG = '#e6edf3' BLUE = '#58a6ff' GREEN = '#3fb950' ORANGE = '#f78166' PURPLE = '#bc8cff' YELLOW = '#e3b341' GRAY = '#8b949e' DARK_GRAY = '#21262d' RED = '#f85149' CYAN = '#39d353' TEAL = '#1abc9c' fig, ax = plt.subplots(figsize=(10, 10), dpi=108) fig.patch.set_facecolor(BG) ax.set_facecolor(BG) ax.set_xlim(0, 10) ax.set_ylim(0, 10) ax.axis('off') TOTAL_FRAMES = 600 PHASE1_END = 150 PHASE2_END = 300 PHASE3_END = 460 PHASE4_END = 600 np.random.seed(13) artists = [] def clear_artists(): global artists for a in artists: try: a.remove() except: pass artists = [] def add(a): artists.append(a) return a def ease(t): return t * t * (3 - 2 * t) def smooth(frame, start, end): if frame <= start: return 0.0 if frame >= end: return 1.0 return ease((frame - start) / (end - start)) def animate(frame): clear_artists() # === PHASE 1: Model-free RL — trial and error (0-150) === if frame < PHASE1_END: p = smooth(frame, 0, 30) add(ax.text(5, 9.4, 'Model-Free RL: метод тыка', color=FG, fontsize=18, ha='center', va='center', fontweight='bold', alpha=p)) # Show random trajectories (failed attempts) np.random.seed(frame // 5) n_trajs = int(smooth(frame, 20, 120) * 8) for i in range(n_trajs): np.random.seed(i * 7 + 3) t = np.linspace(0, 1, 30) traj_x = 1.5 + 7.0 * t + np.cumsum(np.random.randn(30) * 0.15) traj_y = 5.0 + np.cumsum(np.random.randn(30) * 0.3) traj_x = np.clip(traj_x, 0.5, 9.5) traj_y = np.clip(traj_y, 1.0, 9.0) col = RED if i < n_trajs - 1 else GREEN alpha = 0.3 if i < n_trajs - 1 else 0.9 lw = 1.2 if i < n_trajs - 1 else 2.5 line = ax.plot(traj_x, traj_y, color=col, lw=lw, alpha=alpha, zorder=4)[0] add(line) # Start s_a = smooth(frame, 15, 50) s_c = plt.Circle((1.5, 5.0), 0.2, color=BLUE, alpha=s_a, zorder=7) add(ax.add_patch(s_c)) add(ax.text(1.5, 4.4, 'Start', color=BLUE, fontsize=10, ha='center', alpha=s_a)) # Goal g_a = smooth(frame, 20, 60) g_c = plt.Circle((8.5, 5.0), 0.2, color=YELLOW, alpha=g_a, zorder=7) add(ax.add_patch(g_c)) add(ax.text(8.5, 4.4, 'Goal', color=YELLOW, fontsize=10, ha='center', alpha=g_a)) # Label label_a = smooth(frame, 90, 145) add(ax.text(5, 2.8, 'Model-free: без модели мира', color=RED, fontsize=12, ha='center', alpha=label_a, fontweight='bold')) add(ax.text(5, 2.2, 'Миллионы попыток → медленно, дорого', color=RED, fontsize=11, ha='center', alpha=label_a)) # === PHASE 2: RSSM latent space transitions (150-300) === elif frame < PHASE2_END: p = (frame - PHASE1_END) / (PHASE2_END - PHASE1_END) add(ax.text(5, 9.4, 'RSSM: Recurrent State Space Model', color=FG, fontsize=17, ha='center', va='center', fontweight='bold')) add(ax.text(5, 8.9, 'DreamerV3 — Hafner et al., 2023', color=GRAY, fontsize=11, ha='center', va='center')) # Encoder enc_a = smooth(frame, PHASE1_END, PHASE1_END + 50) enc_box = FancyBboxPatch((0.2, 5.5), 2.2, 2.5, boxstyle="round,pad=0.1", facecolor=DARK_GRAY, edgecolor=BLUE, lw=2, alpha=enc_a) add(ax.add_patch(enc_box)) add(ax.text(1.3, 7.6, 'Encoder', color=BLUE, fontsize=12, ha='center', fontweight='bold', alpha=enc_a)) add(ax.text(1.3, 7.0, 'CNN / ViT', color=BLUE, fontsize=10, ha='center', alpha=enc_a * 0.8)) add(ax.text(1.3, 6.3, 'o_t → e_t', color=BLUE, fontsize=10, ha='center', alpha=enc_a * 0.7, style='italic')) # Arrow enc → RSSM enc_arr = smooth(frame, PHASE1_END + 40, PHASE1_END + 80) if enc_arr > 0: arr = FancyArrowPatch((2.5, 6.75), (3.4, 6.75), arrowstyle='->', color=FG, lw=1.5, mutation_scale=13, alpha=enc_arr) add(ax.add_patch(arr)) # RSSM core rssm_a = smooth(frame, PHASE1_END + 60, PHASE1_END + 110) rssm_box = FancyBboxPatch((3.3, 4.8), 3.5, 3.5, boxstyle="round,pad=0.12", facecolor='#1a2332', edgecolor=PURPLE, lw=2.5, alpha=rssm_a) add(ax.add_patch(rssm_box)) add(ax.text(5.05, 8.0, 'RSSM', color=PURPLE, fontsize=14, ha='center', fontweight='bold', alpha=rssm_a)) # RSSM internal states h_a = smooth(frame, PHASE1_END + 80, PHASE1_END + 120) h_box = FancyBboxPatch((3.6, 6.2), 1.2, 0.9, boxstyle="round,pad=0.05", facecolor=PURPLE, alpha=h_a * 0.3) add(ax.add_patch(h_box)) add(ax.text(4.2, 6.65, 'h_t\n(det)', color=PURPLE, fontsize=9, ha='center', alpha=h_a, fontweight='bold')) z_box = FancyBboxPatch((5.3, 6.2), 1.2, 0.9, boxstyle="round,pad=0.05", facecolor=ORANGE, alpha=h_a * 0.3) add(ax.add_patch(z_box)) add(ax.text(5.9, 6.65, 'z_t\n(stoch)', color=ORANGE, fontsize=9, ha='center', alpha=h_a, fontweight='bold')) # RSSM dynamics dyn_a = smooth(frame, PHASE1_END + 100, PHASE1_END + 140) add(ax.text(5.05, 5.5, 'Transition: h_{t+1} = f(h_t, z_t, a_t)', color=PURPLE, fontsize=9, ha='center', alpha=dyn_a, style='italic')) add(ax.text(5.05, 5.05, 'Prior: z_{t+1} ~ p(z|h_{t+1})', color=ORANGE, fontsize=9, ha='center', alpha=dyn_a, style='italic')) # Decoder / Reward dec_a = smooth(frame, PHASE1_END + 120, PHASE2_END) dec_box = FancyBboxPatch((7.0, 5.5), 2.5, 2.5, boxstyle="round,pad=0.1", facecolor=DARK_GRAY, edgecolor=GREEN, lw=2, alpha=dec_a) add(ax.add_patch(dec_box)) add(ax.text(8.25, 7.7, 'Heads', color=GREEN, fontsize=12, ha='center', fontweight='bold', alpha=dec_a)) add(ax.text(8.25, 7.2, 'Decoder\no_t ~ p(o|h,z)', color=GREEN, fontsize=9, ha='center', alpha=dec_a * 0.8, style='italic')) add(ax.text(8.25, 6.5, 'Reward\nr_t ~ p(r|h,z)', color=YELLOW, fontsize=9, ha='center', alpha=dec_a * 0.8, style='italic')) add(ax.text(8.25, 5.9, 'Continue\nc_t ~ p(c|h,z)', color=ORANGE, fontsize=9, ha='center', alpha=dec_a * 0.8, style='italic')) if dec_a > 0: arr2 = FancyArrowPatch((6.9, 6.75), (7.0, 6.75), arrowstyle='->', color=FG, lw=1.5, mutation_scale=13, alpha=dec_a) add(ax.add_patch(arr2)) # Latent space visualization lat_a = smooth(frame, PHASE1_END + 130, PHASE2_END) ax_inset = ax.inset_axes([0.02, 0.02, 0.42, 0.42]) ax_inset.set_facecolor(DARK_GRAY) ax_inset.axis('off') ax_inset.set_title('Латентное пространство', color=FG, fontsize=9, pad=3) np.random.seed(42) pts = np.random.randn(20, 2) * 0.8 ax_inset.scatter(pts[:, 0], pts[:, 1], c=PURPLE, s=30, alpha=lat_a * 0.7) # Trajectory in latent space t_lat = np.linspace(0, 2*np.pi, 10) lx = np.cos(t_lat) * 0.9 ly = np.sin(t_lat) * 0.5 ax_inset.plot(lx, ly, color=TEAL, lw=2, alpha=lat_a) artists.append(ax_inset) # === PHASE 3: Imagination rollouts (300-460) === elif frame < PHASE3_END: p = (frame - PHASE2_END) / (PHASE3_END - PHASE2_END) add(ax.text(5, 9.4, 'Планирование в воображении', color=FG, fontsize=18, ha='center', va='center', fontweight='bold')) # Brain/imagination zone brain_a = smooth(frame, PHASE2_END, PHASE2_END + 50) brain_circle = plt.Circle((5, 6.5), 2.8, color='#1a2332', alpha=brain_a * 0.6, zorder=3) add(ax.add_patch(brain_circle)) brain_border = plt.Circle((5, 6.5), 2.8, color=PURPLE, fill=False, lw=2.5, alpha=brain_a, zorder=4) add(ax.add_patch(brain_border)) add(ax.text(5, 9.0, 'Воображение (Latent Space)', color=PURPLE, fontsize=12, ha='center', alpha=brain_a, fontweight='bold')) # Multiple imagination rollouts n_rollouts = 5 rollout_prog = smooth(frame, PHASE2_END + 40, PHASE3_END - 20) angles_start = np.linspace(np.pi * 0.6, np.pi * 1.4, n_rollouts) for i, ang in enumerate(angles_start): rollout_delay = i * 0.08 r_prog = max(0, min(1, (rollout_prog - rollout_delay) / (1 - rollout_delay))) if r_prog <= 0: continue # Spiral rollout from center t_r = np.linspace(0, r_prog * 2.0, 30) rx = 5.0 + t_r * 0.9 * np.cos(ang + t_r * 0.3) ry = 6.5 + t_r * 0.9 * np.sin(ang + t_r * 0.3) # Color by reward (some paths better) reward_val = [GREEN, TEAL, ORANGE, RED, BLUE][i] line = ax.plot(rx, ry, color=reward_val, lw=2, alpha=0.7, zorder=5)[0] add(line) # End point ep = plt.Circle((rx[-1], ry[-1]), 0.12, color=reward_val, alpha=0.9, zorder=6) add(ax.add_patch(ep)) # Best path highlighted best_a = smooth(frame, PHASE2_END + 140, PHASE3_END) t_best = np.linspace(0, 1.8, 30) bx = 5.0 + t_best * 0.9 * np.cos(angles_start[0] + t_best * 0.3) by = 6.5 + t_best * 0.9 * np.sin(angles_start[0] + t_best * 0.3) add(ax.plot(bx, by, color=GREEN, lw=4, alpha=best_a * 0.9, zorder=7)[0]) add(ax.text(5.5, 4.5, 'Лучший план выбран', color=GREEN, fontsize=12, ha='center', alpha=best_a, fontweight='bold')) # Sample efficiency box se_a = smooth(frame, PHASE2_END + 120, PHASE3_END) se_box = FancyBboxPatch((0.3, 1.5), 9.4, 1.6, boxstyle="round,pad=0.1", facecolor=DARK_GRAY, edgecolor=YELLOW, lw=1.5, alpha=se_a * 0.8) add(ax.add_patch(se_box)) add(ax.text(5, 2.95, 'DreamerV3: обучает actor-critic в латентном пространстве', color=YELLOW, fontsize=11, ha='center', alpha=se_a, fontweight='bold')) add(ax.text(5, 2.35, '15 задач разных доменов — Atari, DMC, Minecraft, робототехника — одна модель', color=YELLOW, fontsize=10, ha='center', alpha=se_a * 0.85)) # === PHASE 4: Real robot using imagined plan (460-600) === else: p = (frame - PHASE3_END) / (PHASE4_END - PHASE3_END) add(ax.text(5, 9.4, 'От воображения к действию', color=FG, fontsize=19, ha='center', va='center', fontweight='bold')) # Three-column: Observation → World Model → Action col_a = smooth(frame, PHASE3_END, PHASE3_END + 50) # Observation obs_box = FancyBboxPatch((0.2, 3.5), 2.8, 5.0, boxstyle="round,pad=0.12", facecolor=DARK_GRAY, edgecolor=BLUE, lw=2, alpha=col_a) add(ax.add_patch(obs_box)) add(ax.text(1.6, 8.2, 'Наблюдение', color=BLUE, fontsize=12, ha='center', fontweight='bold', alpha=col_a)) # Mini image grid for ii in range(3): for jj in range(3): c = [BLUE, GREEN, PURPLE, ORANGE, TEAL, BLUE, YELLOW, RED, GREEN][ii*3+jj] rect = FancyBboxPatch((0.4 + jj * 0.75, 5.8 + ii * 0.75), 0.65, 0.65, boxstyle="square,pad=0", facecolor=c, alpha=col_a * 0.35) add(ax.add_patch(rect)) add(ax.text(1.6, 5.4, 'RGB кадры\n+ состояния', color=BLUE, fontsize=9.5, ha='center', alpha=col_a * 0.8)) # Arrow arr1_a = smooth(frame, PHASE3_END + 40, PHASE3_END + 80) if arr1_a > 0: arr = FancyArrowPatch((3.1, 6.0), (3.9, 6.0), arrowstyle='->', color=FG, lw=2, mutation_scale=16, alpha=arr1_a) add(ax.add_patch(arr)) # World Model wm_a = smooth(frame, PHASE3_END + 60, PHASE3_END + 110) wm_box = FancyBboxPatch((3.8, 3.5), 2.8, 5.0, boxstyle="round,pad=0.12", facecolor='#1a2332', edgecolor=PURPLE, lw=2.5, alpha=wm_a) add(ax.add_patch(wm_box)) add(ax.text(5.2, 8.2, 'World Model', color=PURPLE, fontsize=12, ha='center', fontweight='bold', alpha=wm_a)) # Latent trajectory visualization in world model t_traj = np.linspace(0, 2, 20) lx = 4.2 + t_traj * 0.5 * np.cos(t_traj) ly = 5.8 + t_traj * 0.3 * np.sin(t_traj * 2) lx = np.clip(lx, 4.0, 6.5) ly = np.clip(ly, 4.2, 7.8) wm_prog = smooth(frame, PHASE3_END + 80, PHASE3_END + 130) n_pts = max(2, int(wm_prog * 20)) add(ax.plot(lx[:n_pts], ly[:n_pts], color=TEAL, lw=2, alpha=wm_a * 0.8, zorder=6)[0]) for xx, yy in zip(lx[:n_pts:4], ly[:n_pts:4]): pt = plt.Circle((xx, yy), 0.08, color=PURPLE, alpha=wm_a * 0.7, zorder=7) add(ax.add_patch(pt)) add(ax.text(5.2, 4.5, 'RSSM\nactor-critic\nв латенте', color=PURPLE, fontsize=9.5, ha='center', alpha=wm_a * 0.8)) # Arrow arr2_a = smooth(frame, PHASE3_END + 100, PHASE3_END + 140) if arr2_a > 0: arr2 = FancyArrowPatch((6.7, 6.0), (7.4, 6.0), arrowstyle='->', color=FG, lw=2, mutation_scale=16, alpha=arr2_a) add(ax.add_patch(arr2)) # Robot action act_a = smooth(frame, PHASE3_END + 120, PHASE3_END + 160) act_box = FancyBboxPatch((7.3, 3.5), 2.5, 5.0, boxstyle="round,pad=0.12", facecolor='#1a2e1a', edgecolor=GREEN, lw=2, alpha=act_a) add(ax.add_patch(act_box)) add(ax.text(8.55, 8.2, 'Робот', color=GREEN, fontsize=12, ha='center', fontweight='bold', alpha=act_a)) add(ax.text(8.55, 7.5, 'Действие\nиз плана', color=GREEN, fontsize=10, ha='center', alpha=act_a * 0.8)) # Key metrics metrics_a = smooth(frame, PHASE3_END + 150, PHASE4_END) add(ax.text(5, 3.0, 'DreamerV3: 1 модель — 15 доменов без изменений', color=YELLOW, fontsize=12, ha='center', alpha=metrics_a, fontweight='bold')) add(ax.text(5, 2.35, 'Sample efficiency: в 10x лучше model-free', color=TEAL, fontsize=11, ha='center', alpha=metrics_a)) add(ax.text(5, 1.7, 'Связь: Серия 7 (MCTS) + Серия 16 (диффузия)', color=GRAY, fontsize=10, ha='center', alpha=metrics_a * 0.85)) output_path = '/root/Strategy/content/generated/robot_world_models_animation.mp4' anim = FuncAnimation(fig, animate, frames=TOTAL_FRAMES, interval=1000/24, blit=False) writer = FFMpegWriter(fps=24, bitrate=2000) anim.save(output_path, writer=writer) plt.close(fig) print(f"Saved: {output_path}")