multiwalker_worldmodel/
├─ README.md  ← 本ファイル
├─ pyproject.toml or requirements.txt  ← 依存管理どちらか
├─ configs/
│   └─ default.yaml
├─ src/
│  ├─ envs/
│  │   └─ multiwalker_wrapped.py
│  ├─ models/
│  │   ├─ rssm.py
│  │   ├─ encoders.py
│  │   ├─ decoders.py
│  │   └─ heads.py
│  ├─ agents/
│  │   ├─ dreamer_shared.py
│  │   └─ ppo_shared.py
│  ├─ utils/
│  │   ├─ replay_buffer.py
│  │   ├─ video.py
│  │   ├─ tsne.py
│  │   └─ plotting.py
│  ├─ train_dreamer.py
│  ├─ train_baseline_ppo.py
│  ├─ evaluate.py
│  └─ rollout_record.py
└─ scripts/
   ├─ train_dreamer.sh
   ├─ train_baseline_ppo.sh
   ├─ eval_video.sh
   ├─ eval_latent_tsne.sh
   └─ eval_plots.sh

Create file src/envs/multiwalker_wrapped.py with:
- A PettingZoo Multiwalker wrapper that exposes a Gymnasium-like vectorized API using SuperSuit.
- Support n_agents=3, episode_len=500, add_agent_id=True to append one-hot agent id to observation.
- Implement make_env(seed, vector_envs, normalize_obs, frame_stack).
- Provide reset(seed), step(action_dict or stacked actions), observe() returning stacked obs per agent, and a bridge to single shared policy (concatenate per-agent obs into batch).
- Ensure team reward handling and done truncation consistent with PettingZoo.

Create file src/models/encoders.py with:
- MLPEncoder(in_dim, hidden: List[int], out_dim) -> latent features for RSSM input.

Create file src/models/decoders.py with:
- MLPDecoder(in_dim, hidden: List[int], out_dim) -> reconstruct observation.
- RewardHead(in_dim) -> scalar reward prediction.
- DoneHead(in_dim) -> Bernoulli or logits for terminal prediction.

Create file src/models/heads.py with:
- ActorHead(in_dim, act_dim) outputs action distribution params (DiagGaussian for continuous; Categorical for discrete).
- ValueHead(in_dim) outputs state value.

Create file src/models/rssm.py with:
- RSSM with deterministic GRU h_t and stochastic z_t.
- Functions: init_state(batch), observe(enc_obs_seq) -> posteriors, imagine(policy, horizon) -> latent rollout.
- KL loss calc supporting "free nats" optional.

Create file src/agents/dreamer_shared.py with:
- DreamerSharedAgent holding: Encoder, RSSM, Decoder/Reward/Done heads, ActorHead, ValueHead.
- ReplayBuffer for sequences (import from utils).
- train_step() that:
  1) samples sequences (B, T, obs, act, rew, done)
  2) updates world model via L_world
  3) runs RSSM.imagine() with current actor to produce imagined trajectories of length horizon
  4) computes lambda-returns on imagined reward and updates actor/value
- act(obs_batch) returns actions for all 3 agents using shared policy.
- save/load checkpoints to .pt.

Create file src/agents/ppo_shared.py with:
- Shared-parameter PPO for multi-agent:
  - Single policy network conditioned on agent_id one-hot.
  - Collects trajectories across agents into one buffer.
  - Standard PPO update (clip, value loss, entropy).
- Compatible with Gymnasium-like vector env wrapper.

Create file src/utils/replay_buffer.py with:
- Sequence-capable replay for multi-agent shared-policy.
- API: add(transition per agent), sample(batch_size, seq_len), save/load.

Create file src/utils/video.py with:
- record_episode_video(env, policy, path_mp4 or gif) using imageio-ffmpeg/moviepy.

Create file src/utils/tsne.py with:
- collect_latents(agent, eval_episodes) to dump npz of {h, z} sequences.
- run_tsne(npz) -> 2D embedding, save png.

Create file src/utils/plotting.py with:
- plot_action_distribution(log_dir) over time by dimension.
- plot_value_curves(log_dir) for predicted V vs returns.

Create file src/train_dreamer.py with:
- Load configs/default.yaml.
- Set seeds.
- Build env via envs/multiwalker_wrapped.make_env.
- Instantiate DreamerSharedAgent on device.
- Online loop:
  - init_random_steps: take random actions to fill replay.
  - Main: interact -> store -> if step % update_every == 0: train_step()
  - Periodic eval: run N episodes, log metrics, save video if configured.
  - Save checkpoints at intervals.

Create file src/train_baseline_ppo.py with:
- Similar structure, but using PPOShared agent.
- Train for baseline.steps.

Create file src/evaluate.py with:
- Load checkpoint of Dreamer or PPO.
- Run episodes, compute metrics, optionally save video.
- Option to dump latents and call tsne/plotting utilities.

Create file src/rollout_record.py with:
- Simple tool to roll out a fixed policy and save transitions to npz for analysis.

# Observe (posterior)
h0, z0 = init_state(B)
for t in 1..T:
  e_t = enc(x_t)
  h_t = GRU(h_{t-1}, concat(z_{t-1}, e_t))
  z_t ~ q(z | h_t, e_t)             # posterior
  losses += recon(dec(h_t, z_t), x_t) + β * KL(q || p(z|h_t))
  r_hat = RewardHead(h_t, z_t)
  d_hat = DoneHead(h_t, z_t)

# Imagine (prior)
h0, z0 from last posterior
for τ in 1..H:
  a_τ ~ π(a | h_{τ-1}, z_{τ-1})
  h_τ = GRU(h_{τ-1}, concat(z_{τ-1}, a_τ))
  z_τ ~ p(z | h_τ)                   # prior
  r_τ = RewardHead(h_τ, z_τ)
  d_τ = DoneHead(h_τ, z_τ)
compute λ-returns on {r_τ, d_τ} → targets
update actor/value on imagined trajectory

obs_per_agent = [obs_i]  # i=1..3
obs_i_aug = concat(obs_i, one_hot(agent_id=i))
batch = stack(obs_i_aug over i)
actions = actor(batch)
step env with per-agent actions