File size: 20,055 Bytes
f17a465 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 | <!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>
|