diff --git a/README.md b/README.md
index bed289d..bee83bb 100644
--- a/README.md
+++ b/README.md
@@ -3,7 +3,8 @@ Terra Baselines provides a set of tools to train and evaluate RL policies on the
 
 ## Features
 - Train on multiple devices using PPO with `train.py` (based on [XLand-MiniGrid](https://github.com/corl-team/xland-minigrid))
-- Generate metrics for your checkpoint with `eval.py`
+- **[Experimental]** AlphaZero-style training with MCTS self-play using `train_mcts_alphaZero.py`
+- Generate metrics for your checkpoint with `eval.py` (supports MCTS at inference)
 - Visualize rollouts of your checkpoint with `visualize.py`
 - Run a grid search on the hyperparameters with `train_sweep.py` (orchestrated with [wandb](https://wandb.ai/))
 
@@ -98,12 +99,119 @@ wandb agent $SWEEP_ID &
 wait
 ```
 
-## Eval
+## AlphaZero-Style Training (Experimental)
+
+Train a policy using AlphaZero-style learning with MCTS self-play using `train_mcts_alphaZero.py`. This approach uses Monte Carlo Tree Search during data collection to generate improved policy targets.
+
+### How It Works
+
+1. **Self-Play with MCTS**: The agent plays episodes using MCTS to select actions. MCTS explores the game tree and produces action probabilities based on visit counts.
+2. **Policy Distillation**: The neural network is trained to match the MCTS policy (cross-entropy loss) and predict returns (MSE loss).
+3. **Iterative Improvement**: As the policy improves, MCTS produces better targets, creating a virtuous cycle.
+
+### Run Training
+
+```bash
+DATASET_PATH=/path/to/dataset DATASET_SIZE=1000 python train_mcts_alphaZero.py
+```
+
+### Configuration
+
+Key parameters in `TrainConfig`:
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `num_envs` | 256 | Number of parallel environments |
+| `num_simulations` | 64 | MCTS simulations per action |
+| `max_episode_steps` | 300 | Maximum steps per episode |
+| `batch_size` | 256 | Training batch size |
+| `gamma` | 0.99 | Discount factor |
+| `policy_lr` | 3e-4 | Learning rate |
+
+### Current Status & Known Issues
+
+> **Warning:** This implementation is experimental and may not achieve stable learning yet.
+
+**Potential improvements needed:**
+
+1. **Replay Buffer** - Currently trains only on the most recent rollout data. AlphaZero typically uses a replay buffer (500K-1M samples) to maintain diversity and prevent catastrophic forgetting.
+
+2. **Reduced Training Steps** - Training on all 76,800 samples (300 steps × 256 envs) per iteration may cause overfitting. Consider limiting to ~32 batches per iteration.
+
+3. **Value Normalization** - Returns can vary in scale. Normalizing returns before training may help balance policy and value losses.
+
+4. **Exploration Decay** - The `gumbel_scale` parameter controls exploration. Consider decreasing it over training.
+
+### Example Fix: Adding a Replay Buffer
+
+```python
+from collections import deque
+import numpy as np
+
+class ReplayBuffer:
+    def __init__(self, max_size=500_000):
+        self.buffer = deque(maxlen=max_size)
+    
+    def add_batch(self, obs_dict, prev_actions, pi, returns):
+        batch_size = returns.shape[0]
+        for i in range(batch_size):
+            self.buffer.append({
+                'obs': jax.tree.map(lambda x: x[i], obs_dict),
+                'prev_actions': prev_actions[i],
+                'pi': pi[i],
+                'returns': returns[i]
+            })
+    
+    def sample(self, batch_size):
+        indices = np.random.choice(len(self.buffer), batch_size, replace=False)
+        batch = [self.buffer[i] for i in indices]
+        return {
+            'obs': jax.tree.map(lambda *xs: jnp.stack(xs), *[b['obs'] for b in batch]),
+            'prev_actions': jnp.stack([b['prev_actions'] for b in batch]),
+            'pi': jnp.stack([b['pi'] for b in batch]),
+            'returns': jnp.stack([b['returns'] for b in batch])
+        }
+```
+
+### Monitoring Training
+
+Training logs to wandb with metrics:
+- `train/policy_loss` - Cross-entropy between network policy and MCTS policy
+- `train/value_loss` - MSE between predicted and actual returns
+- `eval/reward` - Average reward during evaluation
+- `eval/terminations` - Number of successful task completions
+
+A healthy training run should show decreasing policy loss as the network learns to match MCTS outputs.
+
+---
+
+## Eval with Monte Carlo Tree Search
 Evaluate your checkpoint with standard metrics using
 ```
 DATASET_PATH=/path/to/dataset DATASET_SIZE=<num_maps_per_type> python eval.py -run <checkpoint_path> -n <num_environments> -steps <num_steps>
 ```
 
+### Options
+| Flag | Description |
+|------|-------------|
+| `--no-mcts` | Use greedy PPO policy instead of MCTS |
+| `--debug` | Print detailed action comparisons for first 20 steps |
+| `-sim <N>` | Number of MCTS simulations (default: 32) |
+
+> ⚠️ **Dtype Note:** MCTS requires specific dtype handling for JAX compatibility. The `fix_env_cfg_dtypes()` function converts agent config fields to `int8` to prevent type promotion issues during tree search.
+
+Example with MCTS planning:
+```
+DATASET_PATH=/path/to/dataset DATASET_SIZE=1000 python eval.py -run checkpoints/tracked-dense.pkl -n 32 -steps 300
+```
+
+Example with greedy PPO (faster, no tree search):
+```
+DATASET_PATH=/path/to/dataset DATASET_SIZE=1000 python eval.py -run checkpoints/tracked-dense.pkl -n 32 -steps 300 --no-mcts
+```
+
+> **Note:** `eval_legacy.py` contains the original evaluation script without MCTS support.
+
 ## Visualize
 Visualize the rollout of your policy with
 ```
@@ -142,7 +250,7 @@ This generates multi-panel plots showing:
 - Terrain change values and action map evolution
 - Combined overlays of all modifications
 
-## Baselines
+## Baselines (Without MCTS at inference)
 We train 2 models capable of solving both foundation and trench type of environments. They differentiate themselves based on the type of agent (wheeled or tracked), and the type of curriculum used to train them (dense reward with single level, or sparse reward with curriculum). All models are trained on 64x64 maps and are stored in the `checkpoints/` folder.
 
 | Checkpoint           | Map Type  | $C_r$ | $S_p$ | $S_w$ | $Coverage$ |
diff --git a/eval.py b/eval.py
index a025893..d2a608e 100644
--- a/eval.py
+++ b/eval.py
@@ -1,293 +1,506 @@
-import numpy as np
-import jax
-import math
-from utils.models import load_neural_network
-from utils.helpers import load_pkl_object
-from terra.env import TerraEnvBatch
-from terra.actions import (
-    WheeledAction,
-    TrackedAction,
-    WheeledActionType,
-    TrackedActionType,
-)
-import jax.numpy as jnp
-from utils.utils_ppo import obs_to_model_input, wrap_action
-
-# from utils.curriculum import Curriculum
-from tensorflow_probability.substrates import jax as tfp
-from train import TrainConfig  # needed for unpickling checkpoints
-
-
-def _append_to_obs(o, obs_log):
-    if obs_log == {}:
-        return {k: v[:, None] for k, v in o.items()}
-    obs_log = {
-        k: jnp.concatenate((v, o[k][:, None]), axis=1) for k, v in obs_log.items()
-    }
-    return obs_log
-
-
-def rollout_episode(
-    env: TerraEnvBatch,
-    model,
-    model_params,
-    env_cfgs,
-    rl_config,
-    max_frames,
-    deterministic,
-    seed,
-):
-    """
-    NOTE: this function assumes it's a tracked agent in the way it computes the stats.
-    """
-    print(f"Using {seed=}")
-    rng = jax.random.PRNGKey(seed)
-    rng, _rng = jax.random.split(rng)
-    rng_reset = jax.random.split(_rng, rl_config.num_test_rollouts)
-    timestep = env.reset(env_cfgs, rng_reset)
-    prev_actions = jnp.zeros(
-        (rl_config.num_test_rollouts, rl_config.num_prev_actions),
-        dtype=jnp.int32
-    )
-
-    tile_size = env_cfgs.tile_size[0].item()
-    move_tiles = env_cfgs.agent.move_tiles[0].item()
-
-    action_type = env.batch_cfg.action_type
-    if action_type == TrackedAction:
-        move_actions = (TrackedActionType.FORWARD, TrackedActionType.BACKWARD)
-        l_actions = ()
-        do_action = TrackedActionType.DO
-    elif action_type == WheeledAction:
-        move_actions = (WheeledActionType.FORWARD, WheeledActionType.BACKWARD)
-        l_actions = (WheeledActionType.WHEELS_LEFT, WheeledActionType.WHEELS_RIGHT)
-        do_action = WheeledActionType.DO
-    else:
-        raise (ValueError(f"{action_type=}"))
-
-    obs = timestep.observation
-    areas = (obs["target_map"] == -1).sum(
-        tuple([i for i in range(len(obs["target_map"].shape))][1:])
-    ) * (tile_size**2)
-    target_maps_init = obs["target_map"].copy()
-    dig_tiles_per_target_map_init = (target_maps_init == -1).sum(
-        tuple([i for i in range(len(target_maps_init.shape))][1:])
-    )
-
-    t_counter = 0
-    reward_seq = []
-    episode_done_once = None
-    episode_length = None
-    move_cumsum = None
-    do_cumsum = None
-    obs_seq = {}
-    while True:
-        obs_seq = _append_to_obs(obs, obs_seq)
-        rng, rng_act, rng_step = jax.random.split(rng, 3)
-        if model is not None:
-            obs_model = obs_to_model_input(timestep.observation, prev_actions, rl_config)
-            v, logits_pi = model.apply(model_params, obs_model)
-            if deterministic:
-                action = np.argmax(logits_pi, axis=-1)
-            else:
-                pi = tfp.distributions.Categorical(logits=logits_pi)
-                action = pi.sample(seed=rng_act)
-            prev_actions = jnp.roll(prev_actions, shift=1, axis=1)
-            prev_actions = prev_actions.at[:, 0].set(action)
-        else:
-            raise RuntimeError("Model is None!")
-        rng_step = jax.random.split(rng_step, rl_config.num_test_rollouts)
-        timestep = env.step(
-            timestep, wrap_action(action, env.batch_cfg.action_type), rng_step
-        )
-        reward = timestep.reward
-        next_obs = timestep.observation
-        done = timestep.info["task_done"]
-
-        reward_seq.append(reward)
-        print(t_counter)
-        print(10 * "=")
-        t_counter += 1
-        if jnp.all(done).item() or t_counter == max_frames:
-            break
-        obs = next_obs
-
-        # Log stats
-        if episode_done_once is None:
-            episode_done_once = done
-        if episode_length is None:
-            episode_length = jnp.zeros_like(done, dtype=jnp.int32)
-        if move_cumsum is None:
-            move_cumsum = jnp.zeros_like(done, dtype=jnp.int32)
-        if do_cumsum is None:
-            do_cumsum = jnp.zeros_like(done, dtype=jnp.int32)
-
-        episode_done_once = episode_done_once | done
-
-        episode_length += ~episode_done_once
-
-        move_cumsum_tmp = jnp.zeros_like(done, dtype=jnp.int32)
-        for move_action in move_actions:
-            move_mask = (action == move_action) * (~episode_done_once)
-            move_cumsum_tmp += move_tiles * tile_size * move_mask.astype(jnp.int32)
-        for l_action in l_actions:
-            l_mask = (action == l_action) * (~episode_done_once)
-            move_cumsum_tmp += 2 * move_tiles * tile_size * l_mask.astype(jnp.int32)
-        move_cumsum += move_cumsum_tmp
-
-        do_cumsum += (action == do_action) * (~episode_done_once)
-
-    # Path efficiency -- only include finished envs
-    move_cumsum *= episode_done_once
-    path_efficiency = (move_cumsum / jnp.sqrt(areas))[episode_done_once]
-    path_efficiency_std = path_efficiency.std()
-    path_efficiency_mean = path_efficiency.mean()
-
-    # Workspaces efficiency -- only include finished envs
-    reference_workspace_area = 0.5 * np.pi * (8**2)
-    n_dig_actions = do_cumsum // 2
-    workspaces_efficiency = (
-        reference_workspace_area
-        * ((n_dig_actions * episode_done_once) / areas)[episode_done_once]
-    )
-    workspaces_efficiency_mean = workspaces_efficiency.mean()
-    workspaces_efficiency_std = workspaces_efficiency.std()
-
-    # Coverage scores
-    dug_tiles_per_action_map = (obs["action_map"] == -1).sum(
-        tuple([i for i in range(len(obs["action_map"].shape))][1:])
-    )
-    coverage_ratios = dug_tiles_per_action_map / dig_tiles_per_target_map_init
-    coverage_scores = episode_done_once + (~episode_done_once) * coverage_ratios
-    coverage_score_mean = coverage_scores.mean()
-    coverage_score_std = coverage_scores.std()
-
-    stats = {
-        "episode_done_once": episode_done_once,
-        "episode_length": episode_length,
-        "path_efficiency": {
-            "mean": path_efficiency_mean,
-            "std": path_efficiency_std,
-        },
-        "workspaces_efficiency": {
-            "mean": workspaces_efficiency_mean,
-            "std": workspaces_efficiency_std,
-        },
-        "coverage": {
-            "mean": coverage_score_mean,
-            "std": coverage_score_std,
-        },
-    }
-    return np.cumsum(reward_seq), stats, obs_seq
-
-
-def print_stats(
-    stats,
-):
-    episode_done_once = stats["episode_done_once"]
-    episode_length = stats["episode_length"]
-    path_efficiency = stats["path_efficiency"]
-    workspaces_efficiency = stats["workspaces_efficiency"]
-    coverage = stats["coverage"]
-
-    completion_rate = 100 * episode_done_once.sum() / len(episode_done_once)
-
-    print("\nStats:\n")
-    print(f"Completion: {completion_rate:.2f}%")
-    # print(f"First episode length average: {episode_length.mean()}")
-    # print(f"First episode length min: {episode_length.min()}")
-    # print(f"First episode length max: {episode_length.max()}")
-    print(
-        f"Path efficiency: {path_efficiency['mean']:.2f} ({path_efficiency['std']:.2f})"
-    )
-    print(
-        f"Workspaces efficiency: {workspaces_efficiency['mean']:.2f} ({workspaces_efficiency['std']:.2f})"
-    )
-    print(f"Coverage: {coverage['mean']:.2f} ({coverage['std']:.2f})")
-
-
-if __name__ == "__main__":
-    import argparse
-
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "-run",
-        "--run_name",
-        type=str,
-        default="ppo_2023_05_09_10_01_23",
-        help="es/ppo trained agent.",
-    )
-    parser.add_argument(
-        "-env",
-        "--env_name",
-        type=str,
-        default="Terra",
-        help="Environment name.",
-    )
-    parser.add_argument(
-        "-n",
-        "--n_envs",
-        type=int,
-        default=1,
-        help="Number of environments.",
-    )
-    parser.add_argument(
-        "-steps",
-        "--n_steps",
-        type=int,
-        default=10,
-        help="Number of steps.",
-    )
-    parser.add_argument(
-        "-d",
-        "--deterministic",
-        type=int,
-        default=0,
-        help="Deterministic. 0 for stochastic, 1 for deterministic.",
-    )
-    parser.add_argument(
-        "-s",
-        "--seed",
-        type=int,
-        default=0,
-        help="Random seed for the environment.",
-    )
-    args, _ = parser.parse_known_args()
-    n_envs = args.n_envs
-
-    log = load_pkl_object(f"{args.run_name}")
-    config = log["train_config"]
-    # from utils.helpers import load_config
-    # config = load_config("agents/Terra/ppo.yaml", 22333, 33222, 5e-04, True, "")["train_config"]
-
-    config.num_test_rollouts = n_envs
-    config.num_devices = 1
-
-    # curriculum = Curriculum(rl_config=config, n_devices=n_devices)
-    # env_cfgs, dofs_count_dict = curriculum.get_cfgs_eval()
-    env_cfgs = log["env_config"]
-    env_cfgs = jax.tree_map(
-        lambda x: x[0][None, ...].repeat(n_envs, 0), env_cfgs
-    )  # take first config and replicate
-    shuffle_maps = True
-    env = TerraEnvBatch(rendering=False, shuffle_maps=shuffle_maps)
-    config.num_embeddings_agent_min = 60
-
-    model = load_neural_network(config, env)
-    model_params = log["model"]
-    # model_params = jax.tree_map(lambda x: x[0], replicated_params)
-    deterministic = bool(args.deterministic)
-    print(f"\nDeterministic = {deterministic}\n")
-
-    cum_rewards, stats, _ = rollout_episode(
-        env,
-        model,
-        model_params,
-        env_cfgs,
-        config,
-        max_frames=args.n_steps,
-        deterministic=deterministic,
-        seed=args.seed,
-    )
-
-    print_stats(stats)
+import numpy as np
+import jax
+import jax.numpy as jnp
+import jax.random as jrandom
+from terra.env import TerraEnvBatch
+from utils.models import get_model_ready
+from utils.helpers import load_pkl_object
+from utils.utils_ppo import obs_to_model_input, wrap_action, policy
+import mctx
+from functools import partial
+import time
+
+from train import TrainConfig  # needed for unpickling checkpoints
+from tensorflow_probability.substrates import jax as tfp
+
+
+def fix_env_cfg_dtypes(env_cfgs):
+    """
+    Fix the dtypes in env_cfgs to prevent JAX type promotion issues.
+    Python ints in the config cause int32 promotion during JAX operations.
+    We need to ensure config values that are used in JAX ops have int8 dtype.
+    """
+    # Fix agent config integer fields that participate in JAX operations
+    fixed_agent = env_cfgs.agent._replace(
+        angles_base=jnp.int8(env_cfgs.agent.angles_base),
+        angles_cabin=jnp.int8(env_cfgs.agent.angles_cabin),
+        max_wheel_angle=jnp.int8(env_cfgs.agent.max_wheel_angle),
+        move_tiles=jnp.int8(env_cfgs.agent.move_tiles),
+        dig_depth=jnp.int8(env_cfgs.agent.dig_depth),
+        height=jnp.int8(env_cfgs.agent.height),
+        width=jnp.int8(env_cfgs.agent.width),
+    )
+    return env_cfgs._replace(agent=fixed_agent)
+
+
+def load_neural_network(config, env):
+    rng = jax.random.PRNGKey(0)
+    model, _ = get_model_ready(rng, config, env)
+    return model
+
+def root_fn(apply_fn, params, timestep, prev_actions, config):
+    obs = timestep.observation
+    inp = obs_to_model_input(obs, prev_actions, config)
+    value, dist = apply_fn(params, inp)
+    return mctx.RootFnOutput(
+        prior_logits=dist.logits,  # unnormalized action logits
+        value=value[:, 0],         # value of the root state
+        embedding=(timestep, prev_actions),  # embedding = (timestep, prev_actions)
+    )
+
+def make_recurrent_fn(env, apply_fn, config):
+    def recurrent_fn(params, rng, actions, embedding):
+        # embedding is (timestep, prev_actions)
+        timestep, prev_actions = embedding
+        rng, rng_env = jrandom.split(rng)
+        rng_envs = jrandom.split(rng_env, config.num_test_rollouts)
+
+        actions = actions.astype(jnp.int32)
+        terra_actions = wrap_action(actions, env.batch_cfg.action_type)
+        next_timestep = env.step(timestep, terra_actions, rng_envs)
+        next_obs = next_timestep.observation
+
+        # Update prev_actions for the simulation
+        next_prev_actions = jnp.roll(prev_actions, shift=1, axis=1)
+        next_prev_actions = next_prev_actions.at[:, 0].set(actions)
+
+        inp = obs_to_model_input(next_obs, next_prev_actions, config)
+        value, dist = apply_fn(params, inp)
+
+        reward = next_timestep.reward
+        done = next_timestep.done
+        discount = (1.0 - done) * config.gamma
+
+        return mctx.RecurrentFnOutput(
+            reward=reward,
+            discount=discount,
+            prior_logits=dist.logits,
+            value=value[:,0],
+        ), (next_timestep, next_prev_actions)
+    return recurrent_fn
+
+
+def make_mcts_step_fn(model, env, config, use_mcts=True):
+    """Create a JIT-compiled MCTS step function for maximum GPU utilization."""
+    
+    def apply_model(params, inp):
+        val, logits_pi = model.apply(params, inp)
+        pi = tfp.distributions.Categorical(logits=logits_pi)
+        return val, pi
+    
+    recurrent_fn = make_recurrent_fn(env, apply_model, config)
+    
+    @jax.jit
+    def mcts_step(params, rng, timestep, prev_actions):
+        """Single MCTS step - fully JIT compiled."""
+        # Compute root
+        obs = timestep.observation
+        inp = obs_to_model_input(obs, prev_actions, config)
+        value, dist = apply_model(params, inp)
+        
+        # Get greedy PPO action for comparison
+        ppo_action = jnp.argmax(dist.logits, axis=-1)
+        
+        root = mctx.RootFnOutput(
+            prior_logits=dist.logits,
+            value=value[:, 0],
+            embedding=(timestep, prev_actions),
+        )
+        
+        # Run MCTS
+        rng, rng_mcts = jrandom.split(rng)
+        policy_output = mctx.gumbel_muzero_policy(
+            params=params,
+            rng_key=rng_mcts,
+            root=root,
+            recurrent_fn=recurrent_fn,
+            num_simulations=config.num_simulations,
+        )
+        
+        mcts_action = policy_output.action.astype(jnp.int32)
+        
+        # Choose which action to use
+        actions = mcts_action if use_mcts else ppo_action
+        
+        # Step environment
+        rng, rng_step = jrandom.split(rng)
+        rng_steps = jrandom.split(rng_step, config.num_test_rollouts)
+        action_type = env.batch_cfg.action_type
+        next_timestep = env.step(timestep, wrap_action(actions, action_type), rng_steps)
+        
+        # Update prev_actions
+        next_prev_actions = jnp.roll(prev_actions, shift=1, axis=1)
+        next_prev_actions = next_prev_actions.at[:, 0].set(actions)
+        
+        # Return both actions for debugging
+        # Return task_done (successful completion) not just done (could be timeout)
+        task_done = next_timestep.info["task_done"]
+        return rng, next_timestep, next_prev_actions, actions, next_timestep.reward, task_done, ppo_action, mcts_action, value[:, 0]
+    
+    return mcts_step
+
+def _append_to_obs(o, obs_log):
+    if obs_log == {}:
+        return {k: v[:, None] for k, v in o.items()}
+    obs_log = {
+        k: jnp.concatenate((v, o[k][:, None]), axis=1) for k, v in obs_log.items()
+    }
+    return obs_log
+
+def rollout_episode(
+    env: TerraEnvBatch,
+    model,
+    model_params,
+    env_cfgs,
+    rl_config,
+    max_frames,
+    deterministic,
+    seed,
+    use_mcts=True,
+    debug=False,
+):
+    rng = jrandom.PRNGKey(seed)
+    rng, _rng = jrandom.split(rng)
+    rng_reset = jrandom.split(_rng, rl_config.num_test_rollouts)
+    timestep = env.reset(env_cfgs, rng_reset)
+
+    # Store these as JAX arrays to avoid repeated .item() calls
+    tile_size = env_cfgs.tile_size[0]
+    move_tiles = env_cfgs.agent.move_tiles[0]
+    action_type = env.batch_cfg.action_type
+
+    # Determine action types based on agent type
+    from terra.actions import (
+        WheeledAction,
+        TrackedAction,
+        WheeledActionType,
+        TrackedActionType,
+    )
+    if action_type == TrackedAction:
+        move_actions = (TrackedActionType.FORWARD, TrackedActionType.BACKWARD)
+        l_actions = ()
+        do_action = TrackedActionType.DO
+        action_names = ['FWD', 'BWD', 'CLK', 'ACLK', 'CAB_CLK', 'CAB_ACLK', 'DO']
+    elif action_type == WheeledAction:
+        move_actions = (WheeledActionType.FORWARD, WheeledActionType.BACKWARD)
+        l_actions = (WheeledActionType.WHEELS_LEFT, WheeledActionType.WHEELS_RIGHT)
+        do_action = WheeledActionType.DO
+        action_names = ['FWD', 'BWD', 'WHL_L', 'WHL_R', 'CAB_CLK', 'CAB_ACLK', 'DO']
+    else:
+        raise (ValueError(f"{action_type=}"))
+
+    obs = timestep.observation
+    tile_size_float = float(tile_size)
+    move_tiles_float = float(move_tiles)
+    areas = (obs["target_map"] == -1).sum(
+        tuple([i for i in range(len(obs["target_map"].shape))][1:])
+    ) * (tile_size_float**2)
+    target_maps_init = obs["target_map"].copy()
+    dig_tiles_per_target_map_init = (target_maps_init == -1).sum(
+        tuple([i for i in range(len(target_maps_init.shape))][1:])
+    )
+
+    # Create JIT-compiled MCTS step function
+    mcts_step = make_mcts_step_fn(model, env, rl_config, use_mcts=use_mcts)
+    
+    # Initialize prev_actions
+    prev_actions = jnp.zeros(
+        (rl_config.num_test_rollouts, rl_config.num_prev_actions),
+        dtype=jnp.int32
+    )
+
+    # Warmup JIT compilation (first call compiles, subsequent calls are fast)
+    print("Warming up JIT compilation...")
+    start_warmup = time.time()
+    rng, timestep, prev_actions, actions, reward, done, ppo_act, mcts_act, values = mcts_step(
+        model_params, rng, timestep, prev_actions
+    )
+    # Block until warmup is complete
+    jax.block_until_ready(actions)
+    print(f"JIT warmup complete in {time.time() - start_warmup:.2f}s")
+    
+    # Reset for actual run
+    rng = jrandom.PRNGKey(seed)
+    rng, _rng = jrandom.split(rng)
+    rng_reset = jrandom.split(_rng, rl_config.num_test_rollouts)
+    timestep = env.reset(env_cfgs, rng_reset)
+    prev_actions = jnp.zeros(
+        (rl_config.num_test_rollouts, rl_config.num_prev_actions),
+        dtype=jnp.int32
+    )
+
+    t_counter = 0
+    reward_seq = []
+    episode_done_once = None
+    episode_length = None
+    move_cumsum = None
+    do_cumsum = None
+    obs_seq = {}
+    
+    # Action tracking for debugging
+    action_counts = {i: 0 for i in range(7)}
+    mcts_ppo_diff_count = 0
+    
+    mode_str = "MCTS" if use_mcts else "PPO (greedy)"
+    print(f"\nStarting rollout for {max_frames} steps with {rl_config.num_test_rollouts} envs using {mode_str}...")
+    start_time = time.time()
+
+    while True:
+        # Run JIT-compiled MCTS step (all GPU work happens here)
+        rng, timestep, prev_actions, actions, reward, done, ppo_act, mcts_act, values = mcts_step(
+            model_params, rng, timestep, prev_actions
+        )
+
+        reward_seq.append(reward)
+        t_counter += 1
+        
+        # Debug: compare MCTS vs PPO actions
+        if debug and t_counter <= 20:
+            # Print first 20 steps in detail
+            ppo_np = np.array(ppo_act)
+            mcts_np = np.array(mcts_act)
+            actions_np = np.array(actions)
+            values_np = np.array(values)
+            rewards_np = np.array(reward)
+            
+            diff_count = (ppo_np != mcts_np).sum()
+            print(f"\n--- Step {t_counter} ---")
+            print(f"PPO actions:  {ppo_np[:8]}  (showing first 8 envs)")
+            print(f"MCTS actions: {mcts_np[:8]}")
+            print(f"Used actions: {actions_np[:8]}")
+            print(f"Values:       {values_np[:8]}")
+            print(f"Rewards:      {rewards_np[:8]}")
+            print(f"MCTS != PPO:  {diff_count}/{len(ppo_np)}")
+        
+        # Track action distribution
+        for a in np.array(actions):
+            action_counts[int(a)] += 1
+        mcts_ppo_diff_count += int((np.array(ppo_act) != np.array(mcts_act)).sum())
+        
+        # Only check termination every N steps to reduce sync overhead
+        if t_counter >= max_frames:
+            break
+        
+        # Check if all done (this forces a sync, but only once per step)
+        all_done = jnp.all(done)
+        if all_done:
+            break
+
+        # Log stats (keep on GPU, no sync needed)
+        if episode_done_once is None:
+            episode_done_once = done
+            episode_length = jnp.zeros_like(done, dtype=jnp.int32)
+            move_cumsum = jnp.zeros_like(done, dtype=jnp.int32)
+            do_cumsum = jnp.zeros_like(done, dtype=jnp.int32)
+
+        episode_done_once = episode_done_once | done
+        episode_length = episode_length + (~episode_done_once).astype(jnp.int32)
+
+        move_cumsum_tmp = jnp.zeros_like(done, dtype=jnp.int32)
+        for move_action in move_actions:
+            move_mask = (actions == move_action) & (~episode_done_once)
+            move_cumsum_tmp = move_cumsum_tmp + (move_tiles_float * tile_size_float * move_mask).astype(jnp.int32)
+        for la in l_actions:
+            l_mask = (actions == la) & (~episode_done_once)
+            move_cumsum_tmp = move_cumsum_tmp + (2 * move_tiles_float * tile_size_float * l_mask).astype(jnp.int32)
+        move_cumsum = move_cumsum + move_cumsum_tmp
+
+        do_cumsum = do_cumsum + ((actions == do_action) & (~episode_done_once)).astype(jnp.int32)
+        
+        # Print progress every 50 steps (reduced from every step)
+        if t_counter % 50 == 0:
+            elapsed = time.time() - start_time
+            steps_per_sec = t_counter / elapsed
+            print(f"Step {t_counter}/{max_frames} | {steps_per_sec:.1f} steps/sec | "
+                  f"Done: {int(episode_done_once.sum())}/{rl_config.num_test_rollouts}")
+
+    elapsed = time.time() - start_time
+    print(f"\nRollout complete: {t_counter} steps in {elapsed:.2f}s ({t_counter/elapsed:.1f} steps/sec)")
+    
+    # Print action distribution
+    total_actions = sum(action_counts.values())
+    print(f"\nAction distribution ({mode_str}):")
+    for i, name in enumerate(action_names):
+        pct = 100 * action_counts[i] / total_actions if total_actions > 0 else 0
+        print(f"  {name} ({i}): {action_counts[i]:6d} ({pct:5.1f}%)")
+    
+    total_comparisons = t_counter * rl_config.num_test_rollouts
+    print(f"\nMCTS != PPO: {mcts_ppo_diff_count}/{total_comparisons} ({100*mcts_ppo_diff_count/total_comparisons:.1f}%)")
+
+    # Final stats computation (sync to CPU only at the end)
+    obs = timestep.observation
+    if episode_done_once is None:
+        episode_done_once = done
+        episode_length = jnp.zeros_like(done, dtype=jnp.int32)
+        move_cumsum = jnp.zeros_like(done, dtype=jnp.int32)
+        do_cumsum = jnp.zeros_like(done, dtype=jnp.int32)
+    
+    move_cumsum = move_cumsum * episode_done_once
+    path_efficiency = (move_cumsum / jnp.sqrt(areas))[episode_done_once]
+    path_efficiency_std = float(path_efficiency.std())
+    path_efficiency_mean = float(path_efficiency.mean())
+
+    reference_workspace_area = 0.5 * np.pi * (8**2)
+    n_dig_actions = do_cumsum // 2
+    workspaces_efficiency = (
+        reference_workspace_area
+        * ((n_dig_actions * episode_done_once) / areas)[episode_done_once]
+    )
+    workspaces_efficiency_mean = float(workspaces_efficiency.mean())
+    workspaces_efficiency_std = float(workspaces_efficiency.std())
+
+    dug_tiles_per_action_map = (obs["action_map"] == -1).sum(
+        tuple([i for i in range(len(obs["action_map"].shape))][1:])
+    )
+    coverage_ratios = dug_tiles_per_action_map / dig_tiles_per_target_map_init
+    coverage_scores = episode_done_once + (~episode_done_once) * coverage_ratios
+    coverage_score_mean = float(coverage_scores.mean())
+    coverage_score_std = float(coverage_scores.std())
+
+    stats = {
+        "episode_done_once": episode_done_once,
+        "episode_length": episode_length,
+        "path_efficiency": {
+            "mean": path_efficiency_mean,
+            "std": path_efficiency_std,
+        },
+        "workspaces_efficiency": {
+            "mean": workspaces_efficiency_mean,
+            "std": workspaces_efficiency_std,
+        },
+        "coverage": {
+            "mean": coverage_score_mean,
+            "std": coverage_score_std,
+        },
+    }
+    
+    return np.cumsum(np.array(reward_seq)), stats, obs_seq
+
+def print_stats(stats):
+    episode_done_once = stats["episode_done_once"]
+    path_efficiency = stats["path_efficiency"]
+    workspaces_efficiency = stats["workspaces_efficiency"]
+    coverage = stats["coverage"]
+
+    completion_rate = 100 * episode_done_once.sum() / len(episode_done_once)
+    print("\nStats:\n")
+    print(f"Completion: {completion_rate:.2f}%")
+    print(
+        f"Path efficiency: {path_efficiency['mean']:.2f} ({path_efficiency['std']:.2f})"
+    )
+    print(
+        f"Workspaces efficiency: {workspaces_efficiency['mean']:.2f} ({workspaces_efficiency['std']:.2f})"
+    )
+    print(f"Coverage: {coverage['mean']:.2f} ({coverage['std']:.2f})")
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-run",
+        "--run_name",
+        type=str,
+        default="checkpoints/tracked-dense.pkl",
+        help="Path to the checkpoint with the trained model.",
+    )
+    parser.add_argument(
+        "-env",
+        "--env_name",
+        type=str,
+        default="Terra",
+        help="Environment name.",
+    )
+    parser.add_argument(
+        "-n",
+        "--n_envs",
+        type=int,
+        default=32,
+        help="Number of environments.",
+    )
+    parser.add_argument(
+        "-steps",
+        "--n_steps",
+        type=int,
+        default=305,
+        help="Number of steps to run.",
+    )
+    parser.add_argument(
+        "-d",
+        "--deterministic",
+        type=int,
+        default=1,
+        help="Deterministic. 0 for stochastic (not directly relevant since MCTS picks argmax?), 1 for deterministic.",
+    )
+    parser.add_argument(
+        "-s",
+        "--seed",
+        type=int,
+        default=42,
+        help="Random seed for the environment.",
+    )
+    parser.add_argument(
+        "--no-mcts",
+        action="store_true",
+        help="Use greedy PPO instead of MCTS (for comparison).",
+    )
+    parser.add_argument(
+        "--debug",
+        action="store_true",
+        help="Print detailed debug info for first 20 steps.",
+    )
+    parser.add_argument(
+        "-sim",
+        "--num_simulations",
+        type=int,
+        default=32,
+        help="Number of MCTS simulations per step.",
+    )
+    args, _ = parser.parse_known_args()
+    n_envs = args.n_envs
+    use_mcts = not args.no_mcts
+
+    log = load_pkl_object(f"{args.run_name}")
+    config = log["train_config"]
+    config.num_test_rollouts = n_envs
+    config.num_devices = 1
+    # Set MCTS parameters
+    config.num_simulations = args.num_simulations
+    if not hasattr(config, 'gamma'):
+        config.gamma = 0.99
+
+    env_cfgs = log["env_config"]
+    env_cfgs = jax.tree_map(
+        lambda x: x[0][None, ...].repeat(n_envs, 0), env_cfgs
+    )  # replicate for n_envs
+    
+    # Fix config dtypes to prevent JAX type promotion issues during MCTS
+    env_cfgs = fix_env_cfg_dtypes(env_cfgs)
+    
+    shuffle_maps = True  # Match eval.py behavior
+    env = TerraEnvBatch(rendering=False, shuffle_maps=shuffle_maps)
+    config.num_embeddings_agent_min = 60
+
+    model = load_neural_network(config, env)
+    model_params = log["model"]
+    deterministic = bool(args.deterministic)
+    
+    mode_str = "MCTS" if use_mcts else "PPO (greedy)"
+    print(f"\nMode: {mode_str}")
+    print(f"MCTS simulations: {config.num_simulations}")
+    print(f"Gamma: {config.gamma}")
+    print(f"Debug: {args.debug}\n")
+
+    cum_rewards, stats, _ = rollout_episode(
+        env,
+        model,
+        model_params,
+        env_cfgs,
+        config,
+        max_frames=args.n_steps,
+        deterministic=deterministic,
+        seed=args.seed,
+        use_mcts=use_mcts,
+        debug=args.debug,
+    )
+
+    print_stats(stats)
\ No newline at end of file
diff --git a/eval_legacy.py b/eval_legacy.py
new file mode 100644
index 0000000..a025893
--- /dev/null
+++ b/eval_legacy.py
@@ -0,0 +1,293 @@
+import numpy as np
+import jax
+import math
+from utils.models import load_neural_network
+from utils.helpers import load_pkl_object
+from terra.env import TerraEnvBatch
+from terra.actions import (
+    WheeledAction,
+    TrackedAction,
+    WheeledActionType,
+    TrackedActionType,
+)
+import jax.numpy as jnp
+from utils.utils_ppo import obs_to_model_input, wrap_action
+
+# from utils.curriculum import Curriculum
+from tensorflow_probability.substrates import jax as tfp
+from train import TrainConfig  # needed for unpickling checkpoints
+
+
+def _append_to_obs(o, obs_log):
+    if obs_log == {}:
+        return {k: v[:, None] for k, v in o.items()}
+    obs_log = {
+        k: jnp.concatenate((v, o[k][:, None]), axis=1) for k, v in obs_log.items()
+    }
+    return obs_log
+
+
+def rollout_episode(
+    env: TerraEnvBatch,
+    model,
+    model_params,
+    env_cfgs,
+    rl_config,
+    max_frames,
+    deterministic,
+    seed,
+):
+    """
+    NOTE: this function assumes it's a tracked agent in the way it computes the stats.
+    """
+    print(f"Using {seed=}")
+    rng = jax.random.PRNGKey(seed)
+    rng, _rng = jax.random.split(rng)
+    rng_reset = jax.random.split(_rng, rl_config.num_test_rollouts)
+    timestep = env.reset(env_cfgs, rng_reset)
+    prev_actions = jnp.zeros(
+        (rl_config.num_test_rollouts, rl_config.num_prev_actions),
+        dtype=jnp.int32
+    )
+
+    tile_size = env_cfgs.tile_size[0].item()
+    move_tiles = env_cfgs.agent.move_tiles[0].item()
+
+    action_type = env.batch_cfg.action_type
+    if action_type == TrackedAction:
+        move_actions = (TrackedActionType.FORWARD, TrackedActionType.BACKWARD)
+        l_actions = ()
+        do_action = TrackedActionType.DO
+    elif action_type == WheeledAction:
+        move_actions = (WheeledActionType.FORWARD, WheeledActionType.BACKWARD)
+        l_actions = (WheeledActionType.WHEELS_LEFT, WheeledActionType.WHEELS_RIGHT)
+        do_action = WheeledActionType.DO
+    else:
+        raise (ValueError(f"{action_type=}"))
+
+    obs = timestep.observation
+    areas = (obs["target_map"] == -1).sum(
+        tuple([i for i in range(len(obs["target_map"].shape))][1:])
+    ) * (tile_size**2)
+    target_maps_init = obs["target_map"].copy()
+    dig_tiles_per_target_map_init = (target_maps_init == -1).sum(
+        tuple([i for i in range(len(target_maps_init.shape))][1:])
+    )
+
+    t_counter = 0
+    reward_seq = []
+    episode_done_once = None
+    episode_length = None
+    move_cumsum = None
+    do_cumsum = None
+    obs_seq = {}
+    while True:
+        obs_seq = _append_to_obs(obs, obs_seq)
+        rng, rng_act, rng_step = jax.random.split(rng, 3)
+        if model is not None:
+            obs_model = obs_to_model_input(timestep.observation, prev_actions, rl_config)
+            v, logits_pi = model.apply(model_params, obs_model)
+            if deterministic:
+                action = np.argmax(logits_pi, axis=-1)
+            else:
+                pi = tfp.distributions.Categorical(logits=logits_pi)
+                action = pi.sample(seed=rng_act)
+            prev_actions = jnp.roll(prev_actions, shift=1, axis=1)
+            prev_actions = prev_actions.at[:, 0].set(action)
+        else:
+            raise RuntimeError("Model is None!")
+        rng_step = jax.random.split(rng_step, rl_config.num_test_rollouts)
+        timestep = env.step(
+            timestep, wrap_action(action, env.batch_cfg.action_type), rng_step
+        )
+        reward = timestep.reward
+        next_obs = timestep.observation
+        done = timestep.info["task_done"]
+
+        reward_seq.append(reward)
+        print(t_counter)
+        print(10 * "=")
+        t_counter += 1
+        if jnp.all(done).item() or t_counter == max_frames:
+            break
+        obs = next_obs
+
+        # Log stats
+        if episode_done_once is None:
+            episode_done_once = done
+        if episode_length is None:
+            episode_length = jnp.zeros_like(done, dtype=jnp.int32)
+        if move_cumsum is None:
+            move_cumsum = jnp.zeros_like(done, dtype=jnp.int32)
+        if do_cumsum is None:
+            do_cumsum = jnp.zeros_like(done, dtype=jnp.int32)
+
+        episode_done_once = episode_done_once | done
+
+        episode_length += ~episode_done_once
+
+        move_cumsum_tmp = jnp.zeros_like(done, dtype=jnp.int32)
+        for move_action in move_actions:
+            move_mask = (action == move_action) * (~episode_done_once)
+            move_cumsum_tmp += move_tiles * tile_size * move_mask.astype(jnp.int32)
+        for l_action in l_actions:
+            l_mask = (action == l_action) * (~episode_done_once)
+            move_cumsum_tmp += 2 * move_tiles * tile_size * l_mask.astype(jnp.int32)
+        move_cumsum += move_cumsum_tmp
+
+        do_cumsum += (action == do_action) * (~episode_done_once)
+
+    # Path efficiency -- only include finished envs
+    move_cumsum *= episode_done_once
+    path_efficiency = (move_cumsum / jnp.sqrt(areas))[episode_done_once]
+    path_efficiency_std = path_efficiency.std()
+    path_efficiency_mean = path_efficiency.mean()
+
+    # Workspaces efficiency -- only include finished envs
+    reference_workspace_area = 0.5 * np.pi * (8**2)
+    n_dig_actions = do_cumsum // 2
+    workspaces_efficiency = (
+        reference_workspace_area
+        * ((n_dig_actions * episode_done_once) / areas)[episode_done_once]
+    )
+    workspaces_efficiency_mean = workspaces_efficiency.mean()
+    workspaces_efficiency_std = workspaces_efficiency.std()
+
+    # Coverage scores
+    dug_tiles_per_action_map = (obs["action_map"] == -1).sum(
+        tuple([i for i in range(len(obs["action_map"].shape))][1:])
+    )
+    coverage_ratios = dug_tiles_per_action_map / dig_tiles_per_target_map_init
+    coverage_scores = episode_done_once + (~episode_done_once) * coverage_ratios
+    coverage_score_mean = coverage_scores.mean()
+    coverage_score_std = coverage_scores.std()
+
+    stats = {
+        "episode_done_once": episode_done_once,
+        "episode_length": episode_length,
+        "path_efficiency": {
+            "mean": path_efficiency_mean,
+            "std": path_efficiency_std,
+        },
+        "workspaces_efficiency": {
+            "mean": workspaces_efficiency_mean,
+            "std": workspaces_efficiency_std,
+        },
+        "coverage": {
+            "mean": coverage_score_mean,
+            "std": coverage_score_std,
+        },
+    }
+    return np.cumsum(reward_seq), stats, obs_seq
+
+
+def print_stats(
+    stats,
+):
+    episode_done_once = stats["episode_done_once"]
+    episode_length = stats["episode_length"]
+    path_efficiency = stats["path_efficiency"]
+    workspaces_efficiency = stats["workspaces_efficiency"]
+    coverage = stats["coverage"]
+
+    completion_rate = 100 * episode_done_once.sum() / len(episode_done_once)
+
+    print("\nStats:\n")
+    print(f"Completion: {completion_rate:.2f}%")
+    # print(f"First episode length average: {episode_length.mean()}")
+    # print(f"First episode length min: {episode_length.min()}")
+    # print(f"First episode length max: {episode_length.max()}")
+    print(
+        f"Path efficiency: {path_efficiency['mean']:.2f} ({path_efficiency['std']:.2f})"
+    )
+    print(
+        f"Workspaces efficiency: {workspaces_efficiency['mean']:.2f} ({workspaces_efficiency['std']:.2f})"
+    )
+    print(f"Coverage: {coverage['mean']:.2f} ({coverage['std']:.2f})")
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-run",
+        "--run_name",
+        type=str,
+        default="ppo_2023_05_09_10_01_23",
+        help="es/ppo trained agent.",
+    )
+    parser.add_argument(
+        "-env",
+        "--env_name",
+        type=str,
+        default="Terra",
+        help="Environment name.",
+    )
+    parser.add_argument(
+        "-n",
+        "--n_envs",
+        type=int,
+        default=1,
+        help="Number of environments.",
+    )
+    parser.add_argument(
+        "-steps",
+        "--n_steps",
+        type=int,
+        default=10,
+        help="Number of steps.",
+    )
+    parser.add_argument(
+        "-d",
+        "--deterministic",
+        type=int,
+        default=0,
+        help="Deterministic. 0 for stochastic, 1 for deterministic.",
+    )
+    parser.add_argument(
+        "-s",
+        "--seed",
+        type=int,
+        default=0,
+        help="Random seed for the environment.",
+    )
+    args, _ = parser.parse_known_args()
+    n_envs = args.n_envs
+
+    log = load_pkl_object(f"{args.run_name}")
+    config = log["train_config"]
+    # from utils.helpers import load_config
+    # config = load_config("agents/Terra/ppo.yaml", 22333, 33222, 5e-04, True, "")["train_config"]
+
+    config.num_test_rollouts = n_envs
+    config.num_devices = 1
+
+    # curriculum = Curriculum(rl_config=config, n_devices=n_devices)
+    # env_cfgs, dofs_count_dict = curriculum.get_cfgs_eval()
+    env_cfgs = log["env_config"]
+    env_cfgs = jax.tree_map(
+        lambda x: x[0][None, ...].repeat(n_envs, 0), env_cfgs
+    )  # take first config and replicate
+    shuffle_maps = True
+    env = TerraEnvBatch(rendering=False, shuffle_maps=shuffle_maps)
+    config.num_embeddings_agent_min = 60
+
+    model = load_neural_network(config, env)
+    model_params = log["model"]
+    # model_params = jax.tree_map(lambda x: x[0], replicated_params)
+    deterministic = bool(args.deterministic)
+    print(f"\nDeterministic = {deterministic}\n")
+
+    cum_rewards, stats, _ = rollout_episode(
+        env,
+        model,
+        model_params,
+        env_cfgs,
+        config,
+        max_frames=args.n_steps,
+        deterministic=deterministic,
+        seed=args.seed,
+    )
+
+    print_stats(stats)
diff --git a/requirements.txt b/requirements.txt
index abcb761..244e02c 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -6,4 +6,5 @@ flax
 matplotlib
 pygame
 wandb
-tensorflow_probability
\ No newline at end of file
+tensorflow_probability
+mctx
\ No newline at end of file
diff --git a/train_mcts_alphaZero.py b/train_mcts_alphaZero.py
new file mode 100644
index 0000000..d38ddbc
--- /dev/null
+++ b/train_mcts_alphaZero.py
@@ -0,0 +1,588 @@
+# train_alphaZero_jitted.py
+import jax
+import jax.numpy as jnp
+import jax.random as jrandom
+from jax import jit, vmap, lax
+from functools import partial
+import optax
+import wandb
+import random
+import mctx
+from dataclasses import dataclass, asdict
+
+from flax import struct
+from flax.training.train_state import TrainState
+
+from terra.env import TerraEnvBatch
+from terra.config import EnvConfig
+import eval_ppo
+import utils.helpers as helpers
+from utils.utils_ppo import obs_to_model_input, wrap_action
+from utils.models import get_model_ready
+
+
+def fix_env_cfg_dtypes(env_cfgs):
+    """
+    Fix the dtypes in env_cfgs to prevent JAX type promotion issues.
+    Python ints in the config cause int32 promotion during JAX operations.
+    We need to ensure config values that are used in JAX ops have int8 dtype.
+    """
+    # Fix agent config integer fields that participate in JAX operations
+    fixed_agent = env_cfgs.agent._replace(
+        angles_base=jnp.int8(env_cfgs.agent.angles_base),
+        angles_cabin=jnp.int8(env_cfgs.agent.angles_cabin),
+        max_wheel_angle=jnp.int8(env_cfgs.agent.max_wheel_angle),
+        move_tiles=jnp.int8(env_cfgs.agent.move_tiles),
+        dig_depth=jnp.int8(env_cfgs.agent.dig_depth),
+        height=jnp.int8(env_cfgs.agent.height),
+        width=jnp.int8(env_cfgs.agent.width),
+    )
+    return env_cfgs._replace(agent=fixed_agent)
+
+
+@dataclass(frozen=True)
+class TrainConfig:
+    name: str = "alphazero-fixed-v1"
+    project: str = "terra-alphazero"
+    group: str = "default"
+
+    num_devices: int = 1
+    num_envs_per_device: int = 256
+    num_envs: int = num_devices * num_envs_per_device
+    max_episode_steps: int = 300
+    episodes_per_iteration: int = 1
+    num_iterations: int = 10000
+
+    gamma: float = 0.99
+    num_simulations: int = 64
+    value_target: str = "maxq"
+
+    batch_size: int = 256
+    policy_lr: float = 3e-4
+    value_lr: float = 3e-4
+    max_grad_norm: float = 0.5
+
+    seed: int = 42
+
+    log_interval: int = 1
+    eval_interval: int = 5
+    checkpoint_interval: int = 10
+
+    total_timesteps: int = 30_000_000_000
+    clip_action_maps: bool = True
+    mask_out_arm_extension: bool = True
+    local_map_normalization_bounds: tuple = (-16, 16)
+    maps_net_normalization_bounds: tuple = (-16, 16)
+    loaded_max: int = 100
+    num_rollouts_eval: int = 300
+    
+    # Required by model
+    num_prev_actions: int = 5
+    num_test_rollouts: int = 32  # For evaluation
+
+    def __getitem__(self, key):
+        return getattr(self, key)
+
+@struct.dataclass
+class SelfPlayTransition:
+    """Stores number of enviroments states for each time step:
+       (obs, pi_mcts, final_return).
+    """
+    obs: jnp.ndarray
+    pi_mcts: jnp.ndarray
+    final_return: jnp.ndarray
+
+
+def make_recurrent_fn(env: TerraEnvBatch, apply_fn, gamma, config: TrainConfig):
+    """Build root & recurrent functions for mctx using the single network.
+    
+    Key fixes:
+    1. Properly tracks prev_actions through MCTS simulations
+    2. Handles terminal states correctly - zeros value/discount when done
+       to prevent MCTS from simulating across episode boundaries
+    """
+
+    def env_step_fn(env_states, actions, rng):
+        bsize = env_states.done.shape[0]
+        rngs = jrandom.split(rng, bsize)
+        wrapped_acts = wrap_action(actions.astype(jnp.int32), env.batch_cfg.action_type)
+        return env.step(env_states, wrapped_acts, rngs)
+
+    def root_fn(params, env_states, prev_actions):
+        """Create root node for MCTS. Embedding = (timestep, prev_actions)."""
+        obs_inp = obs_to_model_input(env_states.observation, prev_actions, config)
+        v, logits = apply_fn(params, obs_inp)
+        return mctx.RootFnOutput(
+            prior_logits=logits,
+            value=v[:, 0],
+            embedding=(env_states, prev_actions)
+        )
+
+    def recurrent_fn(params, rng_key, actions, embedding):
+        """MCTS transition function with proper terminal state handling."""
+        env_states, prev_actions = embedding
+        
+        # Check if already done BEFORE stepping (to avoid simulating from reset states)
+        was_done = env_states.done
+        
+        # Step environment
+        next_env_states = env_step_fn(env_states, actions, rng_key)
+        
+        # Update prev_actions for the simulation
+        next_prev_actions = jnp.roll(prev_actions, shift=1, axis=1)
+        next_prev_actions = next_prev_actions.at[:, 0].set(actions)
+        
+        # Compute network outputs
+        obs_inp = obs_to_model_input(next_env_states.observation, next_prev_actions, config)
+        v, logits = apply_fn(params, obs_inp)
+        
+        reward = next_env_states.reward
+        done = next_env_states.done
+        
+        # KEY FIX: Handle terminal states properly
+        # If was_done: we're simulating from a reset state (invalid), zero everything
+        # If done: episode just ended, zero future value/discount
+        terminal_mask = done | was_done
+        
+        # Zero reward for already-done states (those were invalid transitions)
+        safe_reward = jnp.where(was_done, 0.0, reward)
+        
+        # Zero discount for terminal states (no future value)
+        discount = jnp.where(terminal_mask, 0.0, gamma)
+        
+        # Zero value for terminal states
+        safe_value = jnp.where(terminal_mask, 0.0, v[:, 0])
+        
+        # Mask logits for terminal states to prevent further expansion
+        # (set to very negative so softmax gives ~uniform, but discount=0 prevents use)
+        masked_logits = jnp.where(
+            terminal_mask[:, None],
+            jnp.full_like(logits, -1e9),
+            logits
+        )
+        
+        return mctx.RecurrentFnOutput(
+            reward=safe_reward,
+            discount=discount,
+            prior_logits=masked_logits,
+            value=safe_value
+        ), (next_env_states, next_prev_actions)
+
+    return root_fn, recurrent_fn
+
+
+@partial(jit, static_argnames=("num_simulations", "env", "root_fn", "recurrent_fn", "config"))
+def one_mcts_step(rng, env, states, prev_actions, params,
+                  root_fn, recurrent_fn,
+                  num_simulations,
+                  config):
+    """
+    A single MCTS-guided environment step for all envs in 'states'.
+
+    1) MCTS => action, pi_mcts
+    2) Env step => next state
+    3) Update prev_actions
+    4) Return (next_states, next_prev_actions, actions, pi_mcts).
+    """
+    B = states.done.shape[0]
+    num_actions = 7  # Terra has 7 actions for tracked agent
+    
+    # Invalid actions mask (if needed - currently none masked)
+    invalid_mask = jnp.zeros((B, num_actions), dtype=jnp.bool_)
+
+    rng, rng_mcts, rng_step = jrandom.split(rng, 3)
+    
+    # Create root with prev_actions
+    root = root_fn(params, states, prev_actions)
+    
+    # Run MCTS
+    out = mctx.gumbel_muzero_policy(
+        params=params,
+        rng_key=rng_mcts,
+        root=root,
+        recurrent_fn=recurrent_fn,
+        num_simulations=num_simulations,
+        invalid_actions=invalid_mask,
+        gumbel_scale=1.0,
+    )
+    actions = out.action            # shape [B]
+    pi_mcts = out.action_weights    # shape [B, num_actions]
+
+    # Step environment
+    rngs_step = jrandom.split(rng_step, B)
+    wrapped = wrap_action(actions.astype(jnp.int32), env.batch_cfg.action_type)
+    next_states = env.step(states, wrapped, rngs_step)
+    
+    # Update prev_actions
+    next_prev_actions = jnp.roll(prev_actions, shift=1, axis=1)
+    next_prev_actions = next_prev_actions.at[:, 0].set(actions)
+
+    return rng, next_states, next_prev_actions, actions, pi_mcts
+
+
+def build_obs_buffer_template(obs_dict, max_steps):
+    """
+    For each key in obs_dict, which has shape [B, ...],
+    create a buffer of shape [max_steps, B, ...].
+    """
+    buf_dict = {}
+    for k, arr in obs_dict.items():
+        # arr.shape might be (B, *some_dims)
+        buf_shape = (max_steps,) + arr.shape
+        buf_dict[k] = jnp.zeros(buf_shape, dtype=arr.dtype)
+    return buf_dict
+
+@jit
+def discount_cumsum(rewards, dones, gamma):
+    """
+    Given:
+      rewards: float32[time, batch]
+      dones:   bool[time, batch]
+      gamma:   scalar discount factor
+    Return:
+      returns: float32[time, batch]
+        Where returns[t, b] = rewards[t, b] + gamma * returns[t+1, b], if not done[t].
+        If done[t], then returns[t] = rewards[t] only (no future accumulation).
+    """
+    T, B = rewards.shape
+
+    def scan_fun(carry, t):
+        # t goes from T-1 down to 0
+        future_return = carry  # shape [B]
+        r_t = rewards[t]       # shape [B]
+        done_t = dones[t]      # shape [B]
+        # If done at step t, we do not add gamma * future_return
+        ret_t = r_t + gamma * future_return * (1.0 - done_t.astype(jnp.float32))
+        return ret_t, ret_t
+
+    init = jnp.zeros((B,), dtype=jnp.float32)
+    # We'll scan backwards: range(T-1, ..., 0).
+    # 'lax.scan' runs forward on the given sequence, so we reverse indices.
+    indices = jnp.arange(T - 1, -1, -1)
+    final, all_returns_reversed = lax.scan(scan_fun, init, indices)
+
+    # all_returns_reversed has shape [T, B] in reversed time order
+    # we flip back to [0..T-1].
+    all_returns = jnp.flip(all_returns_reversed, axis=0)
+    return all_returns
+
+@partial(jax.jit, static_argnames=('env', 'root_fn', 'recurrent_fn', 'config'))
+def collect_episodes_jitted(
+    rng,
+    env,             # TerraEnvBatch
+    env_params,      # EnvConfig, repeated for B envs
+    params,          # model params
+    root_fn,         # MCTS root function
+    recurrent_fn,    # MCTS recurrent function
+    config,
+    states,
+    prev_actions,    # [B, num_prev_actions]
+):
+    """Collect self-play data using MCTS, keeping everything on GPU."""
+    B = config.num_envs
+    max_steps = config.max_episode_steps
+    num_actions = 7  # Terra tracked agent
+
+    # Prepare buffers for observations, policies, prev_actions, rewards, dones
+    obs_buf = build_obs_buffer_template(states.observation, max_steps)
+    pi_buf = jnp.zeros((max_steps, B, num_actions), dtype=jnp.float32)
+    prev_actions_buf = jnp.zeros((max_steps, B, config.num_prev_actions), dtype=jnp.int32)
+    reward_buf = jnp.zeros((max_steps, B), dtype=jnp.float32)
+    done_buf = jnp.zeros((max_steps, B), dtype=jnp.bool_)
+
+    step0 = jnp.array(0, dtype=jnp.int32)
+    init_carry = (states, prev_actions, step0, rng, obs_buf, pi_buf, prev_actions_buf, reward_buf, done_buf)
+
+    def cond_fun(carry):
+        (states, prev_actions, step, rng, obs_buf, pi_buf, prev_actions_buf, reward_buf, done_buf) = carry
+        return step < max_steps
+
+    def body_fun(carry):
+        (states, prev_actions, step, rng, obs_buf, pi_buf, prev_actions_buf, reward_buf, done_buf) = carry
+        
+        # Run MCTS step
+        rng, next_states, next_prev_actions, actions, pi_mcts = one_mcts_step(
+            rng, env, states, prev_actions, params,
+            root_fn, recurrent_fn, config.num_simulations, config
+        )
+
+        # Store current obs
+        obs_buf = jax.tree_map(
+            lambda buf, val: buf.at[step].set(val),
+            obs_buf,
+            states.observation
+        )
+        # Store MCTS policy
+        pi_buf = pi_buf.at[step].set(pi_mcts)
+        # Store prev_actions (needed for training)
+        prev_actions_buf = prev_actions_buf.at[step].set(prev_actions)
+        # Store reward & done from NEXT state
+        reward_buf = reward_buf.at[step].set(next_states.reward)
+        done_buf = done_buf.at[step].set(next_states.done)
+
+        return (next_states, next_prev_actions, step + 1, rng, obs_buf, pi_buf, prev_actions_buf, reward_buf, done_buf)
+
+    final_carry = lax.while_loop(cond_fun, body_fun, init_carry)
+    (states_final, prev_actions_final, step_final, rng_final, 
+     obs_buf_final, pi_buf_final, prev_actions_buf_final, reward_buf_final, done_buf_final) = final_carry
+
+    # Compute discounted returns
+    returns_buf = discount_cumsum(reward_buf_final, done_buf_final, config.gamma)
+
+    return (
+        obs_buf_final,          # dict-of-arrays [max_steps, B, ...]
+        pi_buf_final,           # [max_steps, B, num_actions]
+        prev_actions_buf_final, # [max_steps, B, num_prev_actions]
+        returns_buf,            # [max_steps, B]
+        step_final,
+        rng_final,
+        states_final,
+        prev_actions_final,
+    )
+
+@partial(jit, static_argnames=("apply_fn",))
+def alpha_zero_loss(apply_fn, params, obs_batch, pi_batch, returns_batch):
+    """
+    Cross-entropy( pi_batch, pi_pred ) + MSE(value, returns).
+    obs_batch: a batched dict of arrays => pass to obs_to_model_input(...) before call!
+    """
+    v, logits = apply_fn(params, obs_batch)
+    log_probs = jax.nn.log_softmax(logits, axis=-1)
+    probs = jnp.exp(log_probs)
+    # for i in range(pi_batch.shape[0]):
+    #     # jax.debug.print("pi_batch   {}", pi_batch[i].round(3))
+    #     # jax.debug.print("probs      {}", probs[i].round(3))
+    #     # jax.debug.print("sum probs  {}", probs[i].sum())
+
+    #     pi_max_idx = jnp.argmax(pi_batch[i])
+    #     probs_max_idx = jnp.argmax(probs[i])
+    #     match = pi_max_idx == probs_max_idx
+    #     jax.debug.print("Match? {}", match)
+
+    weight_decay = 0.0
+    if weight_decay > 0:
+        def param_l2(p):
+            return jnp.sum(p**2)
+        l2_sum = jax.tree_util.tree_reduce(
+            lambda acc, x: acc + jnp.sum(x**2),
+            params, 
+            initializer=jnp.float32(0.0)
+        )
+        reg_loss = weight_decay * l2_sum
+    else:
+        reg_loss = 0.0
+
+    pol_loss = optax.softmax_cross_entropy(logits=logits, labels=pi_batch).mean()
+    val_loss = jnp.mean((v[:, 0] - returns_batch)**2)
+
+    total_loss = pol_loss + val_loss + reg_loss
+
+    return total_loss, (pol_loss, val_loss, reg_loss)
+
+
+@partial(jit, static_argnames=("apply_fn",))
+def train_step(train_state: TrainState,
+               batch_obs,
+               batch_pi,
+               batch_returns,
+               apply_fn):
+    def loss_fn(params):
+        loss_val, (p_loss, v_loss, reg_loss) = alpha_zero_loss(
+            apply_fn, params,
+            batch_obs, batch_pi, batch_returns
+        )
+        return loss_val, (p_loss, v_loss, reg_loss)
+
+    (loss_val, (p_loss, v_loss, reg_loss)), grads = jax.value_and_grad(loss_fn, has_aux=True)(
+        train_state.params
+    )
+
+    train_state = train_state.apply_gradients(grads=grads)
+    return train_state, (loss_val, p_loss, v_loss, reg_loss)
+
+
+def train_alphazero(config: TrainConfig):
+    """AlphaZero-style training with MCTS self-play."""
+    import time
+    
+    # wandb
+    run = wandb.init(
+        project=config.project,
+        group=config.group,
+        name=config.name,
+        config=asdict(config),
+        save_code=True
+    )
+
+    # Setup
+    rng = jrandom.PRNGKey(config.seed)
+    env = TerraEnvBatch()
+    env_params = EnvConfig()
+    env_params = jax.tree.map(lambda x: jnp.array(x).repeat(config.num_envs, axis=0), env_params)
+    
+    # Fix dtypes to prevent JAX type promotion issues during MCTS
+    env_params = fix_env_cfg_dtypes(env_params)
+
+    # Build network
+    network, network_params = get_model_ready(rng, config, env)
+    
+    # Optionally load pretrained weights
+    # log = helpers.load_pkl_object("checkpoints/your-checkpoint.pkl")
+    # network_params = log["model_params"]
+
+    tx = optax.chain(
+        optax.clip_by_global_norm(config.max_grad_norm), 
+        optax.adam(config.policy_lr)
+    )
+
+    train_state = TrainState.create(apply_fn=network.apply, params=network_params, tx=tx)
+
+    # Create recurrent functions for MCTS
+    root_fn, rec_fn = make_recurrent_fn(env, network.apply, config.gamma, config)
+
+    B = config.num_envs
+
+    # Initialize environment and prev_actions
+    rng, rng_reset = jrandom.split(rng)
+    rng_keys = jrandom.split(rng_reset, B)
+    states = env.reset(env_params, rng_keys)
+    prev_actions = jnp.zeros((B, config.num_prev_actions), dtype=jnp.int32)
+
+    print(f"Starting AlphaZero training with {B} environments, {config.num_simulations} MCTS sims")
+    
+    for iteration in range(config.num_iterations):
+        start_time = time.time()
+        
+        # A) Collect data via MCTS self-play
+        (
+            obs_buf,            # dict-of-arrays: [max_steps, B, ...]
+            pi_buf,             # [max_steps, B, num_actions]
+            prev_actions_buf,   # [max_steps, B, num_prev_actions]
+            returns_buf,        # [max_steps, B]
+            step_final,
+            rng,
+            states_final,
+            prev_actions_final,
+        ) = collect_episodes_jitted(
+            rng, env, env_params,
+            train_state.params,
+            root_fn, rec_fn, config,
+            states, prev_actions
+        )
+
+        # Continue from where we left off
+        states = states_final
+        prev_actions = prev_actions_final
+        
+        collect_time = time.time() - start_time
+
+        # B) Prepare training data
+        T = config.max_episode_steps  # Use full buffer since we run for max_steps
+
+        # Flatten [T, B] -> [T*B]
+        obs_dict_flat = jax.tree.map(
+            lambda arr: arr.reshape((T * B,) + arr.shape[2:]),
+            obs_buf
+        )
+        pi_buf_flat = pi_buf.reshape((T * B, pi_buf.shape[-1]))
+        prev_actions_flat = prev_actions_buf.reshape((T * B, config.num_prev_actions))
+        returns_buf_flat = returns_buf.reshape((T * B,))
+
+        # C) Shuffle data for better training
+        rng, rng_shuffle = jrandom.split(rng)
+        total_samples = T * B
+        perm = jrandom.permutation(rng_shuffle, total_samples)
+        
+        obs_dict_flat = jax.tree.map(lambda x: x[perm], obs_dict_flat)
+        pi_buf_flat = pi_buf_flat[perm]
+        prev_actions_flat = prev_actions_flat[perm]
+        returns_buf_flat = returns_buf_flat[perm]
+
+        # D) Train in mini-batches
+        train_start = time.time()
+        batch_size = config.batch_size
+        num_batches = total_samples // batch_size
+        losses = []
+
+        for batch_idx in range(num_batches):
+            start_idx = batch_idx * batch_size
+            end_idx = start_idx + batch_size
+            
+            obs_dict_batch = jax.tree.map(lambda x: x[start_idx:end_idx], obs_dict_flat)
+            pi_batch = pi_buf_flat[start_idx:end_idx]
+            prev_actions_batch = prev_actions_flat[start_idx:end_idx]
+            ret_batch = returns_buf_flat[start_idx:end_idx]
+
+            # Convert to model input (with prev_actions)
+            inp = obs_to_model_input(obs_dict_batch, prev_actions_batch, config)
+
+            # Gradient step
+            train_state, (loss_val, pol_l, val_l, reg_l) = train_step(
+                train_state, inp, pi_batch, ret_batch, network.apply
+            )
+            losses.append((loss_val, pol_l, val_l, reg_l))
+
+        train_time = time.time() - train_start
+
+        # Compute mean losses
+        if losses:
+            mean_loss = float(jnp.mean(jnp.array([l[0] for l in losses])))
+            mean_pl = float(jnp.mean(jnp.array([l[1] for l in losses])))
+            mean_vl = float(jnp.mean(jnp.array([l[2] for l in losses])))
+            mean_rl = float(jnp.mean(jnp.array([l[3] for l in losses])))
+        else:
+            mean_loss = mean_pl = mean_vl = mean_rl = 0
+
+        # Log to wandb
+        wandb.log({
+            "iteration": iteration,
+            "train/total_loss": mean_loss,
+            "train/policy_loss": mean_pl,
+            "train/value_loss": mean_vl,
+            "train/regularization_loss": mean_rl,
+            "timing/collect_time": collect_time,
+            "timing/train_time": train_time,
+            "timing/samples_per_sec": total_samples / (collect_time + train_time),
+        }, step=iteration)
+
+        print(f"Iter {iteration}: loss={mean_loss:.4f} (pol={mean_pl:.4f}, val={mean_vl:.4f}) | "
+              f"collect={collect_time:.1f}s, train={train_time:.1f}s")
+
+        # Evaluate
+        if (iteration + 1) % config.eval_interval == 0:
+            eval_stats = eval_ppo.rollout(
+                rng, env, env_params,
+                train_state,
+                config
+            )
+            wandb.log({
+                "eval/reward": float(eval_stats.reward) / config.num_envs,
+                "eval/max_reward": float(eval_stats.max_reward),
+                "eval/min_reward": float(eval_stats.min_reward),
+                "eval/episodes": float(eval_stats.episodes),
+                "eval/terminations": float(eval_stats.terminations),
+                "eval/positive_terminations": float(eval_stats.positive_terminations),
+            }, step=iteration)
+            print(f"  Eval: reward={float(eval_stats.reward)/config.num_envs:.2f}, "
+                  f"terminations={float(eval_stats.terminations)}")
+
+        # Checkpoint
+        if (iteration + 1) % config.checkpoint_interval == 0:
+            ckpt = {
+                "iteration": iteration,
+                "model_params": train_state.params,
+                "model": train_state.params,  # For compatibility with eval.py
+                "train_config": config,
+                "env_config": env_params,
+            }
+            helpers.save_pkl_object(ckpt, f"checkpoints/{config.name}.pkl")
+            print(f"  Saved checkpoint to checkpoints/{config.name}.pkl")
+
+    run.finish()
+    return train_state
+
+
+if __name__ == "__main__":
+    cfg = TrainConfig()
+    final_model = train_alphazero(cfg)
+    print("Done training!")
diff --git a/utils/utils_ppo.py b/utils/utils_ppo.py
index 5256d44..c8d2324 100644
--- a/utils/utils_ppo.py
+++ b/utils/utils_ppo.py
@@ -58,5 +58,8 @@ def select_action_ppo(
 
 
 def wrap_action(action, action_type):
-    action = action_type.new(action[:, None])
+    # Explicitly cast to int8 to avoid dtype mismatch issues during JAX tracing
+    # (IntLowDim in terra is jnp.int8)
+    action = jnp.asarray(action, dtype=jnp.int8)[:, None]
+    action = action_type.new(action)
     return action