| <!doctype html> |
| <html lang="en"> |
| <head> |
| <meta charset="utf-8"> |
| <meta name="viewport" content="width=device-width,initial-scale=1"> |
| <title>World Model Implementation Lab</title> |
| <meta name="description" content="Practical world model architectures and PyTorch examples."> |
| <style> |
| :root{ |
| --bg:#070b14;--panel:#0d1424;--panel2:#101a2f;--line:#22304b; |
| --text:#eef5ff;--muted:#9fb0ca;--accent:#82adff |
| } |
| *{box-sizing:border-box} |
| body{ |
| margin:0;background: |
| radial-gradient(circle at 15% -10%,rgba(70,110,255,.18),transparent 34%), |
| radial-gradient(circle at 90% 0%,rgba(139,92,246,.11),transparent 34%), |
| var(--bg); |
| color:var(--text);font-family:Inter,ui-sans-serif,system-ui,-apple-system,BlinkMacSystemFont,"Segoe UI",sans-serif |
| } |
| .wrap{max-width:1450px;margin:0 auto;padding:28px} |
| .hero,.panel{border:1px solid var(--line);border-radius:22px;background:rgba(13,20,36,.94)} |
| .hero{padding:34px;background:linear-gradient(145deg,rgba(18,31,58,.97),rgba(8,14,27,.97))} |
| .kicker{font-size:.76rem;letter-spacing:.13em;text-transform:uppercase;color:#b8c8df;margin-bottom:12px} |
| h1{font-size:clamp(2.2rem,5vw,4.5rem);line-height:.98;margin:0 0 14px} |
| h2,h3{margin-top:0} |
| .hero p,.lead{max-width:950px;color:var(--muted);line-height:1.7} |
| .flow{display:flex;flex-wrap:wrap;gap:8px;align-items:center;margin-top:20px} |
| .pill{border:1px solid var(--line);background:#09111f;border-radius:999px;padding:8px 11px;font-size:.86rem} |
| .arrow{color:#6f84a7} |
| .grid{display:grid;gap:18px;margin-top:18px} |
| .panel{padding:24px} |
| .controls{display:grid;grid-template-columns:2fr 1fr;gap:10px;margin-bottom:14px} |
| input,select{ |
| width:100%;padding:12px 13px;border:1px solid var(--line);border-radius:12px; |
| background:#08111f;color:var(--text);font:inherit |
| } |
| .cards{display:grid;grid-template-columns:repeat(2,minmax(0,1fr));gap:13px} |
| .card{ |
| border:1px solid var(--line);border-radius:16px;background:#09111f;padding:16px; |
| cursor:pointer;transition:.15s ease |
| } |
| .card:hover,.card.active{border-color:#4d6f9f;background:#0d192b} |
| .level{display:inline-block;border:1px solid #334967;border-radius:999px;padding:4px 7px;font-size:.7rem;color:#c2d3ea;margin-bottom:9px} |
| .card h3{font-size:1rem;margin-bottom:6px} |
| .focus{color:#c2d1e5;font-size:.86rem;margin-bottom:8px} |
| .summary{color:var(--muted);font-size:.87rem;line-height:1.5} |
| .tags{display:flex;gap:5px;flex-wrap:wrap;margin-top:10px} |
| .badge{font-size:.71rem;border:1px solid #31435f;border-radius:999px;padding:4px 7px;color:#bed0e8} |
| .detail{display:grid;grid-template-columns:minmax(0,1fr) minmax(0,1.6fr);gap:16px} |
| .detailBox{border:1px solid var(--line);border-radius:16px;background:#09111f;padding:16px} |
| .meta{color:var(--muted);line-height:1.65;font-size:.9rem} |
| pre{ |
| margin:0;max-height:680px;overflow:auto;background:#050a12;border:1px solid #1c2a40;border-radius:14px;padding:16px; |
| color:#dbe8fa;font-family:ui-monospace,SFMono-Regular,Menlo,Consolas,monospace;font-size:.82rem;line-height:1.5 |
| } |
| .file{font-family:ui-monospace,SFMono-Regular,Menlo,Consolas,monospace;color:#9fc1ff} |
| .note{border-left:3px solid var(--accent);background:#091426;border-radius:10px;padding:13px 15px;color:#bccbe0;line-height:1.6} |
| .footer{text-align:center;color:#8295af;font-size:.86rem;padding:28px 0 8px} |
| .footer a{color:#aac7ff} |
| @media(max-width:900px){.detail{grid-template-columns:1fr}.cards{grid-template-columns:1fr}} |
| @media(max-width:620px){.wrap{padding:14px}.hero,.panel{padding:18px;border-radius:16px}.controls{grid-template-columns:1fr}} |
| </style> |
| </head> |
| <body> |
| <div class="wrap"> |
| <section class="hero"> |
| <div class="kicker">WORLD MODEL · IMPLEMENTATION LAB</div> |
| <h1>World Model Implementation Lab</h1> |
| <p> |
| Explore compact PyTorch reference implementations for the core mechanics of a world model: |
| latent state, learned dynamics, action conditioning, imagined rollouts and planning. |
| </p> |
| <div class="flow"> |
| <span class="pill">Observation</span><span class="arrow">→</span> |
| <span class="pill">Encoder</span><span class="arrow">→</span> |
| <span class="pill">Latent State</span><span class="arrow">→</span> |
| <span class="pill">Dynamics</span><span class="arrow">→</span> |
| <span class="pill">Rollout</span><span class="arrow">→</span> |
| <span class="pill">Planning</span> |
| </div> |
| </section> |
|
|
| <div class="grid"> |
| <section class="panel"> |
| <h2>Implementation Library</h2> |
| <div class="controls"> |
| <input id="search" type="search" placeholder="Search dynamics, planning, actions, training…" /> |
| <select id="level"><option value="">All levels</option></select> |
| </div> |
| <div id="cards" class="cards"></div> |
| </section> |
|
|
| <section class="panel"> |
| <div class="detail"> |
| <div class="detailBox"> |
| <span id="detailLevel" class="level"></span> |
| <h2 id="detailTitle"></h2> |
| <div id="detailFocus" class="focus"></div> |
| <p id="detailSummary" class="meta"></p> |
| <p class="meta"><strong>Repository file</strong><br><span id="detailFile" class="file"></span></p> |
| <div id="detailTags" class="tags"></div> |
| <div class="note" style="margin-top:16px"> |
| These examples are intentionally compact. They are reference implementations for learning and experimentation, |
| not production robot or autonomous-system controllers. |
| </div> |
| </div> |
| <pre><code id="code"></code></pre> |
| </div> |
| </section> |
|
|
| <section class="panel"> |
| <h2>Recommended learning path</h2> |
| <div class="flow"> |
| <span class="pill">1 · Latent Dynamics</span><span class="arrow">→</span> |
| <span class="pill">2 · Action Conditioning</span><span class="arrow">→</span> |
| <span class="pill">3 · Rollouts</span><span class="arrow">→</span> |
| <span class="pill">4 · Planning</span><span class="arrow">→</span> |
| <span class="pill">5 · End-to-end Training</span> |
| </div> |
| <p class="lead" style="margin-top:15px"> |
| Start small. First verify one-step state prediction. Then test repeated rollouts. |
| Add planning only after the learned dynamics are useful for the target environment. |
| </p> |
| </section> |
| </div> |
|
|
| <div class="footer"> |
| World Model · Cooperation, research and ecosystem partnerships: |
| <a href="mailto:agenten@magenta.de">agenten@magenta.de</a> |
| </div> |
| </div> |
|
|
| <script> |
| const DATA = [{"id": "latent", "title": "Latent Dynamics", "level": "Foundation", "focus": "Encoder + action-conditioned latent transition", "file": "examples/latent_dynamics.py", "tags": ["Encoder", "Latent State", "Dynamics"], "summary": "The smallest useful architecture for understanding state representation and learned transitions.", "code": "\"\"\"\nMinimal latent-dynamics world model.\n\nArchitecture:\n observation -> encoder -> latent state\n latent state + action -> dynamics -> predicted next latent state\n\nThis example is intentionally small and educational.\n\"\"\"\n\nimport torch\nfrom torch import nn\n\n\nclass Encoder(nn.Module):\n def __init__(self, obs_dim: int, latent_dim: int):\n super().__init__()\n self.net = nn.Sequential(\n nn.Linear(obs_dim, 128),\n nn.ReLU(),\n nn.Linear(128, latent_dim),\n )\n\n def forward(self, obs: torch.Tensor) -> torch.Tensor:\n return self.net(obs)\n\n\nclass DynamicsModel(nn.Module):\n def __init__(self, latent_dim: int, action_dim: int):\n super().__init__()\n self.net = nn.Sequential(\n nn.Linear(latent_dim + action_dim, 128),\n nn.ReLU(),\n nn.Linear(128, latent_dim),\n )\n\n def forward(\n self,\n latent: torch.Tensor,\n action: torch.Tensor,\n ) -> torch.Tensor:\n x = torch.cat([latent, action], dim=-1)\n delta = self.net(x)\n return latent + delta\n\n\nclass LatentWorldModel(nn.Module):\n def __init__(self, obs_dim: int, action_dim: int, latent_dim: int = 32):\n super().__init__()\n self.encoder = Encoder(obs_dim, latent_dim)\n self.dynamics = DynamicsModel(latent_dim, action_dim)\n\n def predict_next_latent(\n self,\n obs: torch.Tensor,\n action: torch.Tensor,\n ) -> torch.Tensor:\n latent = self.encoder(obs)\n return self.dynamics(latent, action)\n\n\nif __name__ == \"__main__\":\n batch_size = 8\n obs_dim = 16\n action_dim = 4\n\n model = LatentWorldModel(obs_dim, action_dim)\n\n obs = torch.randn(batch_size, obs_dim)\n action = torch.randn(batch_size, action_dim)\n next_obs = torch.randn(batch_size, obs_dim)\n\n predicted_next_latent = model.predict_next_latent(obs, action)\n\n # A simple latent target can be obtained by encoding the true next observation.\n # detach() keeps this minimal example focused on dynamics learning.\n with torch.no_grad():\n target_next_latent = model.encoder(next_obs)\n\n loss = nn.functional.mse_loss(\n predicted_next_latent,\n target_next_latent,\n )\n\n print(\"predicted latent shape:\", predicted_next_latent.shape)\n print(\"latent prediction loss:\", float(loss))\n"}, {"id": "action", "title": "Action-Conditioned Prediction", "level": "Core", "focus": "Predict the next observation from state and action", "file": "examples/action_conditioned_model.py", "tags": ["Actions", "Decoder", "Prediction"], "summary": "Adds explicit action features and reconstructs the predicted future observation.", "code": "\"\"\"\nAction-conditioned world model with reconstruction.\n\nThe model learns:\n observation_t -> latent_t\n latent_t + action_t -> latent_t+1\n latent_t+1 -> predicted observation_t+1\n\"\"\"\n\nimport torch\nfrom torch import nn\n\n\nclass ActionConditionedWorldModel(nn.Module):\n def __init__(\n self,\n obs_dim: int,\n action_dim: int,\n latent_dim: int = 32,\n hidden_dim: int = 128,\n ):\n super().__init__()\n\n self.encoder = nn.Sequential(\n nn.Linear(obs_dim, hidden_dim),\n nn.ReLU(),\n nn.Linear(hidden_dim, latent_dim),\n )\n\n self.action_encoder = nn.Sequential(\n nn.Linear(action_dim, 32),\n nn.ReLU(),\n )\n\n self.dynamics = nn.Sequential(\n nn.Linear(latent_dim + 32, hidden_dim),\n nn.ReLU(),\n nn.Linear(hidden_dim, latent_dim),\n )\n\n self.decoder = nn.Sequential(\n nn.Linear(latent_dim, hidden_dim),\n nn.ReLU(),\n nn.Linear(hidden_dim, obs_dim),\n )\n\n def step(\n self,\n latent: torch.Tensor,\n action: torch.Tensor,\n ) -> torch.Tensor:\n action_features = self.action_encoder(action)\n x = torch.cat([latent, action_features], dim=-1)\n delta = self.dynamics(x)\n return latent + delta\n\n def forward(\n self,\n obs: torch.Tensor,\n action: torch.Tensor,\n ):\n latent = self.encoder(obs)\n next_latent = self.step(latent, action)\n next_obs_pred = self.decoder(next_latent)\n return next_obs_pred, next_latent\n\n\nif __name__ == \"__main__\":\n torch.manual_seed(0)\n\n model = ActionConditionedWorldModel(\n obs_dim=12,\n action_dim=3,\n latent_dim=24,\n )\n\n optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n\n obs_t = torch.randn(32, 12)\n action_t = torch.randn(32, 3)\n obs_t1 = torch.randn(32, 12)\n\n pred_obs_t1, _ = model(obs_t, action_t)\n loss = nn.functional.mse_loss(pred_obs_t1, obs_t1)\n\n optimizer.zero_grad()\n loss.backward()\n optimizer.step()\n\n print(\"training loss:\", float(loss))\n"}, {"id": "planning", "title": "Rollout Planning", "level": "Planning", "focus": "Imagine candidate futures and choose an action", "file": "examples/rollout_planning.py", "tags": ["Rollouts", "Reward", "MPC"], "summary": "A compact Model Predictive Control example using imagined trajectories and predicted reward.", "code": "\"\"\"\nMinimal rollout + candidate-action planning example.\n\nThe planner samples candidate action sequences, rolls each sequence through\nthe learned dynamics model, predicts rewards, and chooses the first action\nfrom the highest-scoring imagined trajectory.\n\nThis is a compact example of Model Predictive Control (MPC).\n\"\"\"\n\nimport torch\nfrom torch import nn\n\n\nclass PlanningWorldModel(nn.Module):\n def __init__(self, latent_dim: int, action_dim: int):\n super().__init__()\n\n self.dynamics = nn.Sequential(\n nn.Linear(latent_dim + action_dim, 128),\n nn.Tanh(),\n nn.Linear(128, latent_dim),\n )\n\n self.reward_head = nn.Sequential(\n nn.Linear(latent_dim, 64),\n nn.ReLU(),\n nn.Linear(64, 1),\n )\n\n def step(self, latent, action):\n delta = self.dynamics(torch.cat([latent, action], dim=-1))\n next_latent = latent + delta\n reward = self.reward_head(next_latent).squeeze(-1)\n return next_latent, reward\n\n def rollout(self, initial_latent, actions):\n \"\"\"\n initial_latent: [batch, latent_dim]\n actions: [batch, horizon, action_dim]\n \"\"\"\n state = initial_latent\n total_reward = torch.zeros(\n state.shape[0],\n device=state.device,\n )\n\n for t in range(actions.shape[1]):\n state, reward = self.step(state, actions[:, t])\n total_reward = total_reward + reward\n\n return state, total_reward\n\n\n@torch.no_grad()\ndef choose_action(\n model,\n current_latent,\n action_dim,\n horizon=5,\n num_candidates=256,\n):\n # Repeat the same current state for all candidate futures.\n start = current_latent.expand(num_candidates, -1)\n\n # Continuous actions in [-1, 1].\n candidate_actions = (\n torch.rand(num_candidates, horizon, action_dim, device=start.device)\n * 2.0\n - 1.0\n )\n\n _, scores = model.rollout(start, candidate_actions)\n best = torch.argmax(scores)\n\n # MPC executes only the first action, then replans after the next observation.\n return candidate_actions[best, 0], scores[best]\n\n\nif __name__ == \"__main__\":\n torch.manual_seed(1)\n\n latent_dim = 16\n action_dim = 3\n\n model = PlanningWorldModel(latent_dim, action_dim)\n current_latent = torch.randn(1, latent_dim)\n\n action, imagined_score = choose_action(\n model,\n current_latent,\n action_dim=action_dim,\n horizon=6,\n num_candidates=512,\n )\n\n print(\"chosen action:\", action)\n print(\"imagined trajectory score:\", float(imagined_score))\n"}, {"id": "toy", "title": "Toy World Model", "level": "End-to-end", "focus": "Train dynamics and run a multi-step imagined rollout", "file": "examples/toy_world_model.py", "tags": ["Training", "Transitions", "Rollout"], "summary": "A readable end-to-end training loop on synthetic state transitions.", "code": "\"\"\"\nEnd-to-end toy world-model training example.\n\nSynthetic environment:\n next_state = state + action_effect + small nonlinear drift\n\nThe model learns this transition function and is then used for a multi-step\nrollout. The purpose is to demonstrate the training loop, not realism.\n\"\"\"\n\nimport torch\nfrom torch import nn\n\n\nSTATE_DIM = 4\nACTION_DIM = 2\n\n\ndef environment_step(state: torch.Tensor, action: torch.Tensor) -> torch.Tensor:\n effect = torch.zeros_like(state)\n effect[:, :ACTION_DIM] = action\n drift = 0.05 * torch.sin(state)\n return state + effect + drift\n\n\nclass ToyWorldModel(nn.Module):\n def __init__(self):\n super().__init__()\n self.model = nn.Sequential(\n nn.Linear(STATE_DIM + ACTION_DIM, 128),\n nn.ReLU(),\n nn.Linear(128, 128),\n nn.ReLU(),\n nn.Linear(128, STATE_DIM),\n )\n\n def forward(self, state, action):\n x = torch.cat([state, action], dim=-1)\n delta = self.model(x)\n return state + delta\n\n\ndef make_batch(batch_size=256):\n state = torch.randn(batch_size, STATE_DIM)\n action = torch.rand(batch_size, ACTION_DIM) * 2.0 - 1.0\n next_state = environment_step(state, action)\n return state, action, next_state\n\n\ndef train(steps=500):\n model = ToyWorldModel()\n optimizer = torch.optim.Adam(model.parameters(), lr=2e-3)\n\n for step in range(steps):\n state, action, next_state = make_batch()\n pred = model(state, action)\n loss = nn.functional.mse_loss(pred, next_state)\n\n optimizer.zero_grad()\n loss.backward()\n optimizer.step()\n\n if step % 100 == 0 or step == steps - 1:\n print(f\"step={step:03d} loss={loss.item():.6f}\")\n\n return model\n\n\n@torch.no_grad()\ndef imagined_rollout(model, start_state, actions):\n states = [start_state]\n state = start_state\n\n for action in actions:\n action = action.unsqueeze(0)\n state = model(state, action)\n states.append(state)\n\n return torch.cat(states, dim=0)\n\n\nif __name__ == \"__main__\":\n torch.manual_seed(7)\n\n model = train()\n\n start = torch.zeros(1, STATE_DIM)\n actions = torch.tensor([\n [0.5, 0.0],\n [0.5, 0.2],\n [0.0, 0.5],\n [-0.4, 0.0],\n ])\n\n rollout = imagined_rollout(model, start, actions)\n\n print(\"\\nimagined states:\")\n print(rollout)\n"}]; |
| const $ = id => document.getElementById(id); |
| let selected = DATA[0]?.id || null; |
| |
| function esc(s){ |
| return String(s ?? "").replace(/[&<>"']/g,c=>({"&":"&","<":"<",">":">",'"':""","'":"'"}[c])); |
| } |
| function uniq(xs){ return [...new Set(xs)].sort(); } |
| |
| function addOptions(){ |
| uniq(DATA.map(x=>x.level)).forEach(v=>{ |
| const o=document.createElement("option");o.value=v;o.textContent=v;$("level").appendChild(o); |
| }); |
| } |
| |
| function renderDetail(item){ |
| if(!item) return; |
| selected=item.id; |
| $("detailLevel").textContent=item.level; |
| $("detailTitle").textContent=item.title; |
| $("detailFocus").textContent=item.focus; |
| $("detailSummary").textContent=item.summary; |
| $("detailFile").textContent=item.file; |
| $("detailTags").innerHTML=item.tags.map(t=>`<span class="badge">${esc(t)}</span>`).join(""); |
| $("code").textContent=item.code; |
| document.querySelectorAll(".card").forEach(c=>c.classList.toggle("active",c.dataset.id===selected)); |
| } |
| |
| function renderCards(){ |
| const q=$("search").value.trim().toLowerCase(); |
| const level=$("level").value; |
| |
| const filtered=DATA.filter(x=>{ |
| const h=[x.title,x.level,x.focus,x.summary,x.file,...x.tags].join(" ").toLowerCase(); |
| if(q && !h.includes(q)) return false; |
| if(level && x.level!==level) return false; |
| return true; |
| }); |
| |
| $("cards").innerHTML=filtered.map(x=>` |
| <article class="card ${x.id===selected?"active":""}" data-id="${esc(x.id)}" tabindex="0"> |
| <span class="level">${esc(x.level)}</span> |
| <h3>${esc(x.title)}</h3> |
| <div class="focus">${esc(x.focus)}</div> |
| <div class="summary">${esc(x.summary)}</div> |
| <div class="tags">${x.tags.map(t=>`<span class="badge">${esc(t)}</span>`).join("")}</div> |
| </article> |
| `).join(""); |
| |
| document.querySelectorAll(".card").forEach(card=>{ |
| const activate=()=>{ |
| const item=DATA.find(x=>x.id===card.dataset.id); |
| renderDetail(item); |
| }; |
| card.addEventListener("click",activate); |
| card.addEventListener("keydown",e=>{ |
| if(e.key==="Enter" || e.key===" "){e.preventDefault();activate();} |
| }); |
| }); |
| |
| if(filtered.length && !filtered.some(x=>x.id===selected)) renderDetail(filtered[0]); |
| } |
| |
| addOptions(); |
| $("search").addEventListener("input",renderCards); |
| $("level").addEventListener("change",renderCards); |
| renderCards(); |
| renderDetail(DATA[0]); |
| </script> |
| </body> |
| </html> |
|
|
