From 4463f3bf8e92cb9c52aae68bcd086a8151d39e77 Mon Sep 17 00:00:00 2001 From: Chameleon01 Date: Thu, 27 Nov 2025 19:12:54 +0100 Subject: [PATCH 1/5] MCTS eval clean --- README.md | 21 +- eval.py | 799 ++++++++++++++++++----------- eval_legacy.py | 293 +++++++++++ train_mcts_alphaZero.py | 600 ++++++++++++++++++++++ train_mcts_alphaZero_noenvreset.py | 526 +++++++++++++++++++ utils/utils_ppo.py | 5 +- 6 files changed, 1949 insertions(+), 295 deletions(-) create mode 100644 eval_legacy.py create mode 100644 train_mcts_alphaZero.py create mode 100644 train_mcts_alphaZero_noenvreset.py diff --git a/README.md b/README.md index bed289d..64e93fa 100644 --- a/README.md +++ b/README.md @@ -98,12 +98,31 @@ wandb agent $SWEEP_ID & wait ``` -## Eval +## Eval with Monte Carlo Tree Search Evaluate your checkpoint with standard metrics using ``` DATASET_PATH=/path/to/dataset DATASET_SIZE= python eval.py -run -n -steps ``` +### Options +| Flag | Description | +|------|-------------| +| `--no-mcts` | Use greedy PPO policy instead of MCTS | +| `--debug` | Print detailed action comparisons for first 20 steps | +| `-sim ` | Number of MCTS simulations (default: 32) | + +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 ``` 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/train_mcts_alphaZero.py b/train_mcts_alphaZero.py new file mode 100644 index 0000000..0edb750 --- /dev/null +++ b/train_mcts_alphaZero.py @@ -0,0 +1,600 @@ +# 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 + + +@dataclass(frozen=True) +class TrainConfig: + name: str = "medium-64sim-all-data-g2.0-continue" + 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 = 32 + policy_lr: float = 7e-5 + value_lr: float = 7e-5 + max_grad_norm: float = 0.1 + + seed: int = 42 + + log_interval: int = 1 + eval_interval: int = 2 + checkpoint_interval: int = 5 + + total_timesteps: int = 30_000_000_000 + clip_action_maps = True + mask_out_arm_extension = True + local_map_normalization_bounds = [-16, 16] + loaded_max = 100 + num_rollouts_eval = 300 + + 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.""" + + 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): + obs_inp = obs_to_model_input(env_states.observation, config) + # the network returns (value, logits) + v, logits = apply_fn(params, obs_inp) + return mctx.RootFnOutput( + prior_logits=logits, + value=v[:, 0], + embedding=env_states + ) + + def recurrent_fn(params, rng_key, actions, embedding): + next_env_states = env_step_fn(embedding, actions, rng_key) + obs_inp = obs_to_model_input(next_env_states.observation, config) + v, logits = apply_fn(params, obs_inp) + reward = next_env_states.reward + done = next_env_states.done + discount= (1.0 - done) * gamma + return mctx.RecurrentFnOutput( + reward=reward, + discount=discount, + prior_logits=logits, + value=v[:, 0] + ), next_env_states + + return root_fn, recurrent_fn + + +@partial(jit, static_argnames=("num_simulations", "env", "root_fn", "recurrent_fn", "config")) +def one_mcts_step(rng, env, states, done_mask, params, + root_fn, recurrent_fn, + num_simulations, + config): + """ + A single environment step for *all* envs in 'states', skipping where done_mask=1. + + 1) MCTS => action, pi_mcts + 2) Env step => next state + 3) Return (next_states, actions, pi_mcts). + """ + # invalid actions mask + B = states.done.shape[0] + invalid_mask = jnp.zeros((B, 9), dtype=jnp.bool_) + invalid_mask = invalid_mask.at[:,-2].set(True).at[:,-3].set(True) + + rng, rng_mcts, rng_step = jrandom.split(rng, 3) + root = root_fn(params, states) + 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, # <--- pass here + gumbel_scale=2.0, + ) + actions = out.action # shape [B] + pi_mcts = out.action_weights # shape [B, num_actions] + + rngs_step = jrandom.split(rng_step, states.done.shape[0]) + wrapped = wrap_action(actions.astype(jnp.int32), env.batch_cfg.action_type) + next_states = env.step(states, wrapped, rngs_step) + + return rng, next_states, 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 +): + B = config.num_envs + max_steps = config.max_episode_steps + + # 1) Reset environment + # rng_keys = jax.random.split(rng, B) + # states = env.reset(env_params, rng_keys) # shape [B] + + # 2) Prepare buffers + obs_buf = build_obs_buffer_template(states.observation, max_steps) + num_actions = 9 + pi_buf = jnp.zeros((max_steps, B, num_actions), dtype=jnp.float32) + + # (NEW) We'll also store the raw reward and done flag at each step + 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, step0, rng, obs_buf, pi_buf, reward_buf, done_buf) + + def cond_fun(carry): + (states, step, rng, obs_buf, pi_buf, reward_buf, done_buf) = carry + all_done = jnp.all(states.done) + return step < max_steps + + def body_fun(carry): + (states, step, rng, obs_buf, pi_buf, reward_buf, done_buf) = carry + rng, next_states, actions, pi_mcts = one_mcts_step( + rng, env, states, states.done, params, + root_fn, recurrent_fn, config.num_simulations, config + ) + + # store current obs in obs_buf + obs_buf = jax.tree_map( + lambda buf, val: buf.at[step].set(val), + obs_buf, + states.observation + ) + # store pi_mcts + pi_buf = pi_buf.at[step].set(pi_mcts) + + # (NEW) store the 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, step + 1, rng, obs_buf, pi_buf, reward_buf, done_buf) + + states_final, step_final, rng_final, obs_buf_final, pi_buf_final, reward_buf_final, done_buf_final = \ + lax.while_loop(cond_fun, body_fun, init_carry) + + ######################################## + # (NEW) Post-processing: compute discounted returns for each step + ######################################## + # step_final tells us how many time steps were actually used (<= max_steps). + # We'll slice the valid portion. + # reward_buf_used = reward_buf_final[:step_final] # shape [step_final, B] + # done_buf_used = done_buf_final[:step_final] # shape [step_final, B] + + # returns_buf has shape [step_final, B]. + returns_buf = discount_cumsum(reward_buf_final, done_buf_final, config.gamma) + + ######################################## + # Return everything we need + ######################################## + return ( + obs_buf_final, # dict-of-arrays [max_steps, B, ...] + pi_buf_final, # [max_steps, B, num_actions] + returns_buf, # (NEW) discounted returns, shape [step_final, B] + step_final, + rng_final, + states_final + ) + +# def collect_episodes( +# rng, +# env: TerraEnvBatch, +# env_params, +# params, +# root_fn, +# recurrent_fn, +# config: TrainConfig +# ): +# """ +# Collect up to config.max_episode_steps for config.num_envs in parallel, then store the data: +# obs[i] => pi[i], for each step in each environment, +# final_return = sum_of_rewards for that environment. +# """ +# rng_keys = jrandom.split(rng, config.num_envs) +# states = env.reset(env_params, rng_keys) # shape [B] + +# B = states.done.shape[0] +# done = states.done +# step = 0 + +# obs_buffers = [[] for _ in range(B)] +# pi_buffers = [[] for _ in range(B)] +# rew_buffers = [[] for _ in range(B)] # <--- store step rewards + +# rng_main = rng + +# while (not bool(jnp.all(done))) and (step < config.max_episode_steps): +# rng_main, next_states, actions, pi_mcts = one_mcts_step( +# rng_main, env, states, done, +# params, +# root_fn, recurrent_fn, +# config.num_simulations, +# config +# ) +# done_np = jax.device_get(done) +# states_obs_np = {} +# for k, arr in states.observation.items(): +# states_obs_np[k] = jax.device_get(arr) +# pi_mcts_np = jax.device_get(pi_mcts) +# rew_np = jax.device_get(next_states.reward) +# for i in range(B): +# if not done_np[i]: +# sub_dict = {kk: states_obs_np[kk][i] for kk in states_obs_np.keys()} +# obs_buffers[i].append(sub_dict) +# pi_buffers[i].append(pi_mcts_np[i]) +# rew_buffers[i].append(float(rew_np[i])) # store step reward as float + +# done = jnp.logical_or(done, next_states.done) +# states = next_states +# step += 1 + +# transitions = [] +# for i in range(B): +# # We have T = len(obs_buffers[i]) steps +# T = len(obs_buffers[i]) +# if T==0: +# continue +# # compute discounted return for each step +# returns = [] +# running = 0.0 +# for t in reversed(range(T)): +# running = rew_buffers[i][t] + config.gamma * running +# returns.append(running) +# returns.reverse() # match forward indices + +# for t in range(len(obs_buffers[i])): +# transitions.append( +# SelfPlayTransition( +# obs=obs_buffers[i][t], +# pi_mcts=jnp.array(pi_buffers[i][t], dtype=jnp.float32), +# final_return=returns[t] +# ) +# ) +# return transitions, rng_main + + +@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): + # 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) + + # Build network using get_model_ready (a single net for both policy & value) + network, network_params = get_model_ready(rng, config, env) + # Preload pretrained weights from checkpoint + log = helpers.load_pkl_object("checkpoints/medium-64sim-all-data-g2.0-debug.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 using the single network + root_fn, rec_fn = make_recurrent_fn(env, network.apply, config.gamma, config) + + + # #evalutate at start + # 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=0) + # print("Eval stats: ", { + # "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), + # }) + + B = config.num_envs + + rng_keys = jax.random.split(rng, B) + states = env.reset(env_params, rng_keys) # shape [B] + + for iteration in range(config.num_iterations): + print("start collect episodes") + # A) Collect data via MCTS self-play + ( + obs_buf, # dict-of-arrays: [max_steps, B, ...] + pi_buf, # [max_steps, B, num_actions] + returns_buf, # (NEW) [step_final, B] + step_final, + rng, + states_final + ) = collect_episodes_jitted( + rng, env, env_params, + train_state.params, + root_fn, rec_fn, config, + states + ) + + states = states_final + + + # B) Flatten or slice time+batch dimension for training + # 'step_final' is how many steps we actually used. + # We slice obs_buf and pi_buf accordingly. + T = step_final # integer <= max_steps + + # slice each obs_dict from [max_steps, B, ...] -> [T, B, ...] + obs_dict_slice = jax.tree_map(lambda arr: arr[:T], obs_buf) + pi_buf_slice = pi_buf[:T] # shape [T, B, num_actions] + + # flatten [T, B] -> [T*B] + obs_dict_flat = jax.tree_map( + lambda arr: arr.reshape((T * B,) + arr.shape[2:]), + obs_dict_slice + ) + pi_buf_flat = pi_buf_slice.reshape((T * B, pi_buf.shape[-1])) + returns_buf_flat = returns_buf.reshape((T * B,)) + + # C) Train in mini-batches + total_samples = T * B + idx = 0 + batch_size = config.batch_size + losses = [] + + while idx < total_samples: + slice_end = idx + batch_size + obs_dict_batch = jax.tree_map(lambda x: x[idx:slice_end], obs_dict_flat) + pi_batch = pi_buf_flat[idx:slice_end] + ret_batch = returns_buf_flat[idx:slice_end] + idx = slice_end + + # convert to model input + inp = obs_to_model_input(obs_dict_batch, config) + + # single step of gradient descent + 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)) + + if losses: + arr_loss = jnp.array([l[0] for l in losses]) + arr_pl = jnp.array([l[1] for l in losses]) + arr_vl = jnp.array([l[2] for l in losses]) + arr_rl = jnp.array([l[3] for l in losses]) + + mean_loss = float(arr_loss.mean()) + mean_pl = float(arr_pl.mean()) + mean_vl = float(arr_vl.mean()) + mean_rl = float(arr_rl.mean()) + else: + mean_loss = 0 + mean_pl = 0 + mean_vl = 0 + mean_rl = 0 + + wandb.log({ + "iteration": iteration, + "train/total_loss": mean_loss, + "train/policy_loss": mean_pl, + "train/value_loss": mean_vl, + "train/regularization_loss": mean_rl + }, step=iteration) + + # Evaluate raw policy + 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) + + # log action counts + wandb.log({ + "action_counts/action_0": float(eval_stats.action_0) / config.num_envs, + "action_counts/action_1": float(eval_stats.action_1) / config.num_envs, + "action_counts/action_2": float(eval_stats.action_2) / config.num_envs, + "action_counts/action_3": float(eval_stats.action_3) / config.num_envs, + "action_counts/action_4": float(eval_stats.action_4) / config.num_envs, + "action_counts/action_5": float(eval_stats.action_5) / config.num_envs, + "action_counts/action_6": float(eval_stats.action_6) / config.num_envs, + "action_counts/action_7": float(eval_stats.action_7) / config.num_envs, + "action_counts/action_8": float(eval_stats.action_8) / config.num_envs, + }, step=iteration) + + # checkpoint + if (iteration+1) % config.checkpoint_interval == 0: + ckpt = { + "iteration": iteration, + "model_params": train_state.params, + "train_config": config, + "env_config": env_params, + } + helpers.save_pkl_object(ckpt, f"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/train_mcts_alphaZero_noenvreset.py b/train_mcts_alphaZero_noenvreset.py new file mode 100644 index 0000000..24dc7d6 --- /dev/null +++ b/train_mcts_alphaZero_noenvreset.py @@ -0,0 +1,526 @@ +# 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 + + +@dataclass(frozen=True) +class TrainConfig: + name: str = "Local-trials-one" + project: str = "terra-alphazero" + group: str = "default" + + num_devices: int = 1 + num_envs_per_device: int = 2 + num_envs: int = num_devices * num_envs_per_device + max_episode_steps: int = 300 + episodes_per_iteration: int = 10 + num_iterations: int = 10000 + + gamma: float = 0.99 + num_simulations: int = 16 + value_target: str = "maxq" + + batch_size: int = 128 + policy_lr: float = 7e-5 + value_lr: float = 7e-5 + max_grad_norm: float = 0.1 + + seed: int = 42 + + log_interval: int = 1 + eval_interval: int = 2 + checkpoint_interval: int = 5 + + total_timesteps: int = 30_000_000_000 + clip_action_maps = True + mask_out_arm_extension = True + local_map_normalization_bounds = [-16, 16] + loaded_max = 100 + num_rollouts_eval = 300 + + 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.""" + + 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): + obs_inp = obs_to_model_input(env_states.observation, config) + # the network returns (value, logits) + v, logits = apply_fn(params, obs_inp) + return mctx.RootFnOutput( + prior_logits=logits, + value=v[:, 0], + embedding=env_states + ) + + def recurrent_fn(params, rng_key, actions, embedding): + next_env_states = env_step_fn(embedding, actions, rng_key) + obs_inp = obs_to_model_input(next_env_states.observation, config) + v, logits = apply_fn(params, obs_inp) + reward = next_env_states.reward + done = next_env_states.done + discount= (1.0 - done) * gamma + return mctx.RecurrentFnOutput( + reward=reward, + discount=discount, + prior_logits=logits, + value=v[:, 0] + ), next_env_states + + return root_fn, recurrent_fn + + +@partial(jit, static_argnames=("num_simulations", "env", "root_fn", "recurrent_fn", "config")) +def one_mcts_step(rng, env, states, done_mask, params, + root_fn, recurrent_fn, + num_simulations, + config): + """ + A single environment step for *all* envs in 'states', skipping where done_mask=1. + + 1) MCTS => action, pi_mcts + 2) Env step => next state + 3) Return (next_states, actions, pi_mcts). + """ + # invalid actions mask + B = states.done.shape[0] + invalid_mask = jnp.zeros((B, 9), dtype=jnp.bool_) + invalid_mask = invalid_mask.at[:,-2].set(True).at[:,-3].set(True) + + rng, rng_mcts, rng_step = jrandom.split(rng, 3) + root = root_fn(params, states) + 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, # <--- pass here + gumbel_scale=2.0, + ) + actions = out.action # shape [B] + pi_mcts = out.action_weights # shape [B, num_actions] + + rngs_step = jrandom.split(rng_step, states.done.shape[0]) + wrapped = wrap_action(actions.astype(jnp.int32), env.batch_cfg.action_type) + next_states = env.step(states, wrapped, rngs_step) + + return rng, next_states, 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 +): + B = config.num_envs + max_steps = config.max_episode_steps + + # 1) Reset environment + # rng_keys = jax.random.split(rng, B) + # states = env.reset(env_params, rng_keys) # shape [B] + + # 2) Prepare buffers + obs_buf = build_obs_buffer_template(states.observation, max_steps) + num_actions = 9 + pi_buf = jnp.zeros((max_steps, B, num_actions), dtype=jnp.float32) + + # (NEW) We'll also store the raw reward and done flag at each step + 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, step0, rng, obs_buf, pi_buf, reward_buf, done_buf) + + def cond_fun(carry): + (states, step, rng, obs_buf, pi_buf, reward_buf, done_buf) = carry + all_done = jnp.all(states.done) + return step < max_steps + + def body_fun(carry): + (states, step, rng, obs_buf, pi_buf, reward_buf, done_buf) = carry + rng, next_states, actions, pi_mcts = one_mcts_step( + rng, env, states, states.done, params, + root_fn, recurrent_fn, config.num_simulations, config + ) + + # store current obs in obs_buf + obs_buf = jax.tree_map( + lambda buf, val: buf.at[step].set(val), + obs_buf, + states.observation + ) + # store pi_mcts + pi_buf = pi_buf.at[step].set(pi_mcts) + + # (NEW) store the 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, step + 1, rng, obs_buf, pi_buf, reward_buf, done_buf) + + states_final, step_final, rng_final, obs_buf_final, pi_buf_final, reward_buf_final, done_buf_final = \ + lax.while_loop(cond_fun, body_fun, init_carry) + + ######################################## + # (NEW) Post-processing: compute discounted returns for each step + ######################################## + # step_final tells us how many time steps were actually used (<= max_steps). + # We'll slice the valid portion. + # reward_buf_used = reward_buf_final[:step_final] # shape [step_final, B] + # done_buf_used = done_buf_final[:step_final] # shape [step_final, B] + + # returns_buf has shape [step_final, B]. + returns_buf = discount_cumsum(reward_buf_final, done_buf_final, config.gamma) + + ######################################## + # Return everything we need + ######################################## + return ( + obs_buf_final, # dict-of-arrays [max_steps, B, ...] + pi_buf_final, # [max_steps, B, num_actions] + returns_buf, # (NEW) discounted returns, shape [step_final, B] + step_final, + rng_final, + states_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): + # 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) + + # Build network using get_model_ready (a single net for both policy & value) + network, network_params = get_model_ready(rng, config, env) + # Preload pretrained weights from checkpoint + # log = helpers.load_pkl_object("checkpoints/medium-64sim-with-epoch.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 using the single network + root_fn, rec_fn = make_recurrent_fn(env, network.apply, config.gamma, config) + + + # #evalutate at start + # 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=0) + # print("Eval stats: ", { + # "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), + # }) + + B = config.num_envs + + rng_keys = jax.random.split(rng, B) + states = env.reset(env_params, rng_keys) # shape [B] + + for iteration in range(config.num_iterations): + print("start collect episodes") + # A) Collect data via MCTS self-play + ( + obs_buf, # dict-of-arrays: [max_steps, B, ...] + pi_buf, # [max_steps, B, num_actions] + returns_buf, # (NEW) [step_final, B] + step_final, + rng, + states_final + ) = collect_episodes_jitted( + rng, env, env_params, + train_state.params, + root_fn, rec_fn, config, + states + ) + + states = states_final + + + # B) Flatten or slice time+batch dimension for training + # 'step_final' is how many steps we actually used. + # We slice obs_buf and pi_buf accordingly. + T = step_final # integer <= max_steps + + # slice each obs_dict from [max_steps, B, ...] -> [T, B, ...] + obs_dict_slice = jax.tree_map(lambda arr: arr[:T], obs_buf) + pi_buf_slice = pi_buf[:T] # shape [T, B, num_actions] + + # flatten [T, B] -> [T*B] + obs_dict_flat = jax.tree_map( + lambda arr: arr.reshape((T * B,) + arr.shape[2:]), + obs_dict_slice + ) + pi_buf_flat = pi_buf_slice.reshape((T * B, pi_buf.shape[-1])) + returns_buf_flat = returns_buf.reshape((T * B,)) + + # C) Train in mini-batches + + for epoch in range(config.episodes_per_iteration): + total_samples = T * B + idx = 0 + batch_size = config.batch_size + losses = [] + + while idx < total_samples: + slice_end = idx + batch_size + obs_dict_batch = jax.tree_map(lambda x: x[idx:slice_end], obs_dict_flat) + pi_batch = pi_buf_flat[idx:slice_end] + ret_batch = returns_buf_flat[idx:slice_end] + idx = slice_end + + # convert to model input + inp = obs_to_model_input(obs_dict_batch, config) + + # single step of gradient descent + 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)) + + if losses: + arr_loss = jnp.array([l[0] for l in losses]) + arr_pl = jnp.array([l[1] for l in losses]) + arr_vl = jnp.array([l[2] for l in losses]) + arr_rl = jnp.array([l[3] for l in losses]) + + mean_loss = float(arr_loss.mean()) + mean_pl = float(arr_pl.mean()) + mean_vl = float(arr_vl.mean()) + mean_rl = float(arr_rl.mean()) + else: + mean_loss = 0 + mean_pl = 0 + mean_vl = 0 + mean_rl = 0 + + wandb.log({ + "iteration": iteration, + "train/total_loss": mean_loss, + "train/policy_loss": mean_pl, + "train/value_loss": mean_vl, + "train/regularization_loss": mean_rl + }, step=iteration) + + # Evaluate raw policy + 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) + + # log action counts + wandb.log({ + "action_counts/action_0": float(eval_stats.action_0) / config.num_envs, + "action_counts/action_1": float(eval_stats.action_1) / config.num_envs, + "action_counts/action_2": float(eval_stats.action_2) / config.num_envs, + "action_counts/action_3": float(eval_stats.action_3) / config.num_envs, + "action_counts/action_4": float(eval_stats.action_4) / config.num_envs, + "action_counts/action_5": float(eval_stats.action_5) / config.num_envs, + "action_counts/action_6": float(eval_stats.action_6) / config.num_envs, + "action_counts/action_7": float(eval_stats.action_7) / config.num_envs, + "action_counts/action_8": float(eval_stats.action_8) / config.num_envs, + }, step=iteration) + + # checkpoint + if (iteration+1) % config.checkpoint_interval == 0: + ckpt = { + "iteration": iteration, + "model_params": train_state.params, + "train_config": config, + "env_config": env_params, + } + helpers.save_pkl_object(ckpt, f"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 From eda9bff58fa1a8e6f262300936f0384d2ea21bcb Mon Sep 17 00:00:00 2001 From: Chameleon01 Date: Thu, 27 Nov 2025 19:34:39 +0100 Subject: [PATCH 2/5] working alpha zero script --- train_mcts_alphaZero.py | 486 +++++++++++++------------- train_mcts_alphaZero_noenvreset.py | 526 ----------------------------- 2 files changed, 237 insertions(+), 775 deletions(-) delete mode 100644 train_mcts_alphaZero_noenvreset.py diff --git a/train_mcts_alphaZero.py b/train_mcts_alphaZero.py index 0edb750..d38ddbc 100644 --- a/train_mcts_alphaZero.py +++ b/train_mcts_alphaZero.py @@ -21,9 +21,28 @@ 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 = "medium-64sim-all-data-g2.0-continue" + name: str = "alphazero-fixed-v1" project: str = "terra-alphazero" group: str = "default" @@ -38,23 +57,28 @@ class TrainConfig: num_simulations: int = 64 value_target: str = "maxq" - batch_size: int = 32 - policy_lr: float = 7e-5 - value_lr: float = 7e-5 - max_grad_norm: float = 0.1 + 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 = 2 - checkpoint_interval: int = 5 + eval_interval: int = 5 + checkpoint_interval: int = 10 total_timesteps: int = 30_000_000_000 - clip_action_maps = True - mask_out_arm_extension = True - local_map_normalization_bounds = [-16, 16] - loaded_max = 100 - num_rollouts_eval = 300 + 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) @@ -70,7 +94,13 @@ class SelfPlayTransition: def make_recurrent_fn(env: TerraEnvBatch, apply_fn, gamma, config: TrainConfig): - """Build root & recurrent functions for mctx using the single network.""" + """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] @@ -78,69 +108,116 @@ def env_step_fn(env_states, actions, rng): 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): - obs_inp = obs_to_model_input(env_states.observation, config) - # the network returns (value, logits) + 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 + embedding=(env_states, prev_actions) ) def recurrent_fn(params, rng_key, actions, embedding): - next_env_states = env_step_fn(embedding, actions, rng_key) - obs_inp = obs_to_model_input(next_env_states.observation, config) + """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 - discount= (1.0 - done) * gamma + + 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=reward, + reward=safe_reward, discount=discount, - prior_logits=logits, - value=v[:, 0] - ), next_env_states + 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, done_mask, params, +def one_mcts_step(rng, env, states, prev_actions, params, root_fn, recurrent_fn, num_simulations, config): """ - A single environment step for *all* envs in 'states', skipping where done_mask=1. + A single MCTS-guided environment step for all envs in 'states'. 1) MCTS => action, pi_mcts 2) Env step => next state - 3) Return (next_states, actions, pi_mcts). + 3) Update prev_actions + 4) Return (next_states, next_prev_actions, actions, pi_mcts). """ - # invalid actions mask B = states.done.shape[0] - invalid_mask = jnp.zeros((B, 9), dtype=jnp.bool_) - invalid_mask = invalid_mask.at[:,-2].set(True).at[:,-3].set(True) + 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) - root = root_fn(params, states) + + # 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, # <--- pass here - gumbel_scale=2.0, + invalid_actions=invalid_mask, + gumbel_scale=1.0, ) actions = out.action # shape [B] pi_mcts = out.action_weights # shape [B, num_actions] - rngs_step = jrandom.split(rng_step, states.done.shape[0]) + # 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, actions, pi_mcts + return rng, next_states, next_prev_actions, actions, pi_mcts def build_obs_buffer_template(obs_dict, max_steps): @@ -198,157 +275,71 @@ def collect_episodes_jitted( root_fn, # MCTS root function recurrent_fn, # MCTS recurrent function config, - states + 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 - # 1) Reset environment - # rng_keys = jax.random.split(rng, B) - # states = env.reset(env_params, rng_keys) # shape [B] - - # 2) Prepare buffers + # Prepare buffers for observations, policies, prev_actions, rewards, dones obs_buf = build_obs_buffer_template(states.observation, max_steps) - num_actions = 9 pi_buf = jnp.zeros((max_steps, B, num_actions), dtype=jnp.float32) - - # (NEW) We'll also store the raw reward and done flag at each step + 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_) + done_buf = jnp.zeros((max_steps, B), dtype=jnp.bool_) step0 = jnp.array(0, dtype=jnp.int32) - init_carry = (states, step0, rng, obs_buf, pi_buf, reward_buf, done_buf) + init_carry = (states, prev_actions, step0, rng, obs_buf, pi_buf, prev_actions_buf, reward_buf, done_buf) def cond_fun(carry): - (states, step, rng, obs_buf, pi_buf, reward_buf, done_buf) = carry - all_done = jnp.all(states.done) + (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, step, rng, obs_buf, pi_buf, reward_buf, done_buf) = carry - rng, next_states, actions, pi_mcts = one_mcts_step( - rng, env, states, states.done, params, + (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 in obs_buf + # Store current obs obs_buf = jax.tree_map( lambda buf, val: buf.at[step].set(val), obs_buf, states.observation ) - # store pi_mcts + # Store MCTS policy pi_buf = pi_buf.at[step].set(pi_mcts) - - # (NEW) store the reward & done from NEXT state + # 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) + done_buf = done_buf.at[step].set(next_states.done) - return (next_states, step + 1, rng, obs_buf, pi_buf, reward_buf, done_buf) + return (next_states, next_prev_actions, step + 1, rng, obs_buf, pi_buf, prev_actions_buf, reward_buf, done_buf) - states_final, step_final, rng_final, obs_buf_final, pi_buf_final, reward_buf_final, done_buf_final = \ - lax.while_loop(cond_fun, body_fun, init_carry) + 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 - ######################################## - # (NEW) Post-processing: compute discounted returns for each step - ######################################## - # step_final tells us how many time steps were actually used (<= max_steps). - # We'll slice the valid portion. - # reward_buf_used = reward_buf_final[:step_final] # shape [step_final, B] - # done_buf_used = done_buf_final[:step_final] # shape [step_final, B] - - # returns_buf has shape [step_final, B]. + # Compute discounted returns returns_buf = discount_cumsum(reward_buf_final, done_buf_final, config.gamma) - ######################################## - # Return everything we need - ######################################## return ( - obs_buf_final, # dict-of-arrays [max_steps, B, ...] - pi_buf_final, # [max_steps, B, num_actions] - returns_buf, # (NEW) discounted returns, shape [step_final, B] + 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 + states_final, + prev_actions_final, ) -# def collect_episodes( -# rng, -# env: TerraEnvBatch, -# env_params, -# params, -# root_fn, -# recurrent_fn, -# config: TrainConfig -# ): -# """ -# Collect up to config.max_episode_steps for config.num_envs in parallel, then store the data: -# obs[i] => pi[i], for each step in each environment, -# final_return = sum_of_rewards for that environment. -# """ -# rng_keys = jrandom.split(rng, config.num_envs) -# states = env.reset(env_params, rng_keys) # shape [B] - -# B = states.done.shape[0] -# done = states.done -# step = 0 - -# obs_buffers = [[] for _ in range(B)] -# pi_buffers = [[] for _ in range(B)] -# rew_buffers = [[] for _ in range(B)] # <--- store step rewards - -# rng_main = rng - -# while (not bool(jnp.all(done))) and (step < config.max_episode_steps): -# rng_main, next_states, actions, pi_mcts = one_mcts_step( -# rng_main, env, states, done, -# params, -# root_fn, recurrent_fn, -# config.num_simulations, -# config -# ) -# done_np = jax.device_get(done) -# states_obs_np = {} -# for k, arr in states.observation.items(): -# states_obs_np[k] = jax.device_get(arr) -# pi_mcts_np = jax.device_get(pi_mcts) -# rew_np = jax.device_get(next_states.reward) -# for i in range(B): -# if not done_np[i]: -# sub_dict = {kk: states_obs_np[kk][i] for kk in states_obs_np.keys()} -# obs_buffers[i].append(sub_dict) -# pi_buffers[i].append(pi_mcts_np[i]) -# rew_buffers[i].append(float(rew_np[i])) # store step reward as float - -# done = jnp.logical_or(done, next_states.done) -# states = next_states -# step += 1 - -# transitions = [] -# for i in range(B): -# # We have T = len(obs_buffers[i]) steps -# T = len(obs_buffers[i]) -# if T==0: -# continue -# # compute discounted return for each step -# returns = [] -# running = 0.0 -# for t in reversed(range(T)): -# running = rew_buffers[i][t] + config.gamma * running -# returns.append(running) -# returns.reverse() # match forward indices - -# for t in range(len(obs_buffers[i])): -# transitions.append( -# SelfPlayTransition( -# obs=obs_buffers[i][t], -# pi_mcts=jnp.array(pi_buffers[i][t], dtype=jnp.float32), -# final_return=returns[t] -# ) -# ) -# return transitions, rng_main - - @partial(jit, static_argnames=("apply_fn",)) def alpha_zero_loss(apply_fn, params, obs_batch, pi_batch, returns_batch): """ @@ -411,6 +402,9 @@ def loss_fn(params): def train_alphazero(config: TrainConfig): + """AlphaZero-style training with MCTS self-play.""" + import time + # wandb run = wandb.init( project=config.project, @@ -424,135 +418,138 @@ def train_alphazero(config: TrainConfig): 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) + 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 using get_model_ready (a single net for both policy & value) + # Build network network, network_params = get_model_ready(rng, config, env) - # Preload pretrained weights from checkpoint - log = helpers.load_pkl_object("checkpoints/medium-64sim-all-data-g2.0-debug.pkl") - network_params = log["model_params"] - - tx = optax.chain(optax.clip_by_global_norm(config.max_grad_norm), optax.adam(config.policy_lr)) + + # 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 using the single network + # Create recurrent functions for MCTS root_fn, rec_fn = make_recurrent_fn(env, network.apply, config.gamma, config) - - # #evalutate at start - # 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=0) - # print("Eval stats: ", { - # "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), - # }) - B = config.num_envs - rng_keys = jax.random.split(rng, B) - states = env.reset(env_params, rng_keys) # shape [B] + # 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): - print("start collect episodes") + 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] - returns_buf, # (NEW) [step_final, B] + 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 + states_final, + prev_actions_final, ) = collect_episodes_jitted( rng, env, env_params, train_state.params, root_fn, rec_fn, config, - states + states, prev_actions ) + # Continue from where we left off states = states_final + prev_actions = prev_actions_final + collect_time = time.time() - start_time - # B) Flatten or slice time+batch dimension for training - # 'step_final' is how many steps we actually used. - # We slice obs_buf and pi_buf accordingly. - T = step_final # integer <= max_steps - - # slice each obs_dict from [max_steps, B, ...] -> [T, B, ...] - obs_dict_slice = jax.tree_map(lambda arr: arr[:T], obs_buf) - pi_buf_slice = pi_buf[:T] # shape [T, B, num_actions] + # 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( + # Flatten [T, B] -> [T*B] + obs_dict_flat = jax.tree.map( lambda arr: arr.reshape((T * B,) + arr.shape[2:]), - obs_dict_slice + obs_buf ) - pi_buf_flat = pi_buf_slice.reshape((T * B, pi_buf.shape[-1])) + 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) Train in mini-batches + # C) Shuffle data for better training + rng, rng_shuffle = jrandom.split(rng) total_samples = T * B - idx = 0 + 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 = [] - while idx < total_samples: - slice_end = idx + batch_size - obs_dict_batch = jax.tree_map(lambda x: x[idx:slice_end], obs_dict_flat) - pi_batch = pi_buf_flat[idx:slice_end] - ret_batch = returns_buf_flat[idx:slice_end] - idx = slice_end + 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 - inp = obs_to_model_input(obs_dict_batch, config) + # Convert to model input (with prev_actions) + inp = obs_to_model_input(obs_dict_batch, prev_actions_batch, config) - # single step of gradient descent + # 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: - arr_loss = jnp.array([l[0] for l in losses]) - arr_pl = jnp.array([l[1] for l in losses]) - arr_vl = jnp.array([l[2] for l in losses]) - arr_rl = jnp.array([l[3] for l in losses]) - - mean_loss = float(arr_loss.mean()) - mean_pl = float(arr_pl.mean()) - mean_vl = float(arr_vl.mean()) - mean_rl = float(arr_rl.mean()) + 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 = 0 - mean_pl = 0 - mean_vl = 0 - mean_rl = 0 + 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 + "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) - # Evaluate raw policy - if (iteration+1) % config.eval_interval == 0: + 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, @@ -566,29 +563,20 @@ def train_alphazero(config: TrainConfig): "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)}") - # log action counts - wandb.log({ - "action_counts/action_0": float(eval_stats.action_0) / config.num_envs, - "action_counts/action_1": float(eval_stats.action_1) / config.num_envs, - "action_counts/action_2": float(eval_stats.action_2) / config.num_envs, - "action_counts/action_3": float(eval_stats.action_3) / config.num_envs, - "action_counts/action_4": float(eval_stats.action_4) / config.num_envs, - "action_counts/action_5": float(eval_stats.action_5) / config.num_envs, - "action_counts/action_6": float(eval_stats.action_6) / config.num_envs, - "action_counts/action_7": float(eval_stats.action_7) / config.num_envs, - "action_counts/action_8": float(eval_stats.action_8) / config.num_envs, - }, step=iteration) - - # checkpoint - if (iteration+1) % config.checkpoint_interval == 0: + # 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 diff --git a/train_mcts_alphaZero_noenvreset.py b/train_mcts_alphaZero_noenvreset.py deleted file mode 100644 index 24dc7d6..0000000 --- a/train_mcts_alphaZero_noenvreset.py +++ /dev/null @@ -1,526 +0,0 @@ -# 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 - - -@dataclass(frozen=True) -class TrainConfig: - name: str = "Local-trials-one" - project: str = "terra-alphazero" - group: str = "default" - - num_devices: int = 1 - num_envs_per_device: int = 2 - num_envs: int = num_devices * num_envs_per_device - max_episode_steps: int = 300 - episodes_per_iteration: int = 10 - num_iterations: int = 10000 - - gamma: float = 0.99 - num_simulations: int = 16 - value_target: str = "maxq" - - batch_size: int = 128 - policy_lr: float = 7e-5 - value_lr: float = 7e-5 - max_grad_norm: float = 0.1 - - seed: int = 42 - - log_interval: int = 1 - eval_interval: int = 2 - checkpoint_interval: int = 5 - - total_timesteps: int = 30_000_000_000 - clip_action_maps = True - mask_out_arm_extension = True - local_map_normalization_bounds = [-16, 16] - loaded_max = 100 - num_rollouts_eval = 300 - - 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.""" - - 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): - obs_inp = obs_to_model_input(env_states.observation, config) - # the network returns (value, logits) - v, logits = apply_fn(params, obs_inp) - return mctx.RootFnOutput( - prior_logits=logits, - value=v[:, 0], - embedding=env_states - ) - - def recurrent_fn(params, rng_key, actions, embedding): - next_env_states = env_step_fn(embedding, actions, rng_key) - obs_inp = obs_to_model_input(next_env_states.observation, config) - v, logits = apply_fn(params, obs_inp) - reward = next_env_states.reward - done = next_env_states.done - discount= (1.0 - done) * gamma - return mctx.RecurrentFnOutput( - reward=reward, - discount=discount, - prior_logits=logits, - value=v[:, 0] - ), next_env_states - - return root_fn, recurrent_fn - - -@partial(jit, static_argnames=("num_simulations", "env", "root_fn", "recurrent_fn", "config")) -def one_mcts_step(rng, env, states, done_mask, params, - root_fn, recurrent_fn, - num_simulations, - config): - """ - A single environment step for *all* envs in 'states', skipping where done_mask=1. - - 1) MCTS => action, pi_mcts - 2) Env step => next state - 3) Return (next_states, actions, pi_mcts). - """ - # invalid actions mask - B = states.done.shape[0] - invalid_mask = jnp.zeros((B, 9), dtype=jnp.bool_) - invalid_mask = invalid_mask.at[:,-2].set(True).at[:,-3].set(True) - - rng, rng_mcts, rng_step = jrandom.split(rng, 3) - root = root_fn(params, states) - 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, # <--- pass here - gumbel_scale=2.0, - ) - actions = out.action # shape [B] - pi_mcts = out.action_weights # shape [B, num_actions] - - rngs_step = jrandom.split(rng_step, states.done.shape[0]) - wrapped = wrap_action(actions.astype(jnp.int32), env.batch_cfg.action_type) - next_states = env.step(states, wrapped, rngs_step) - - return rng, next_states, 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 -): - B = config.num_envs - max_steps = config.max_episode_steps - - # 1) Reset environment - # rng_keys = jax.random.split(rng, B) - # states = env.reset(env_params, rng_keys) # shape [B] - - # 2) Prepare buffers - obs_buf = build_obs_buffer_template(states.observation, max_steps) - num_actions = 9 - pi_buf = jnp.zeros((max_steps, B, num_actions), dtype=jnp.float32) - - # (NEW) We'll also store the raw reward and done flag at each step - 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, step0, rng, obs_buf, pi_buf, reward_buf, done_buf) - - def cond_fun(carry): - (states, step, rng, obs_buf, pi_buf, reward_buf, done_buf) = carry - all_done = jnp.all(states.done) - return step < max_steps - - def body_fun(carry): - (states, step, rng, obs_buf, pi_buf, reward_buf, done_buf) = carry - rng, next_states, actions, pi_mcts = one_mcts_step( - rng, env, states, states.done, params, - root_fn, recurrent_fn, config.num_simulations, config - ) - - # store current obs in obs_buf - obs_buf = jax.tree_map( - lambda buf, val: buf.at[step].set(val), - obs_buf, - states.observation - ) - # store pi_mcts - pi_buf = pi_buf.at[step].set(pi_mcts) - - # (NEW) store the 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, step + 1, rng, obs_buf, pi_buf, reward_buf, done_buf) - - states_final, step_final, rng_final, obs_buf_final, pi_buf_final, reward_buf_final, done_buf_final = \ - lax.while_loop(cond_fun, body_fun, init_carry) - - ######################################## - # (NEW) Post-processing: compute discounted returns for each step - ######################################## - # step_final tells us how many time steps were actually used (<= max_steps). - # We'll slice the valid portion. - # reward_buf_used = reward_buf_final[:step_final] # shape [step_final, B] - # done_buf_used = done_buf_final[:step_final] # shape [step_final, B] - - # returns_buf has shape [step_final, B]. - returns_buf = discount_cumsum(reward_buf_final, done_buf_final, config.gamma) - - ######################################## - # Return everything we need - ######################################## - return ( - obs_buf_final, # dict-of-arrays [max_steps, B, ...] - pi_buf_final, # [max_steps, B, num_actions] - returns_buf, # (NEW) discounted returns, shape [step_final, B] - step_final, - rng_final, - states_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): - # 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) - - # Build network using get_model_ready (a single net for both policy & value) - network, network_params = get_model_ready(rng, config, env) - # Preload pretrained weights from checkpoint - # log = helpers.load_pkl_object("checkpoints/medium-64sim-with-epoch.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 using the single network - root_fn, rec_fn = make_recurrent_fn(env, network.apply, config.gamma, config) - - - # #evalutate at start - # 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=0) - # print("Eval stats: ", { - # "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), - # }) - - B = config.num_envs - - rng_keys = jax.random.split(rng, B) - states = env.reset(env_params, rng_keys) # shape [B] - - for iteration in range(config.num_iterations): - print("start collect episodes") - # A) Collect data via MCTS self-play - ( - obs_buf, # dict-of-arrays: [max_steps, B, ...] - pi_buf, # [max_steps, B, num_actions] - returns_buf, # (NEW) [step_final, B] - step_final, - rng, - states_final - ) = collect_episodes_jitted( - rng, env, env_params, - train_state.params, - root_fn, rec_fn, config, - states - ) - - states = states_final - - - # B) Flatten or slice time+batch dimension for training - # 'step_final' is how many steps we actually used. - # We slice obs_buf and pi_buf accordingly. - T = step_final # integer <= max_steps - - # slice each obs_dict from [max_steps, B, ...] -> [T, B, ...] - obs_dict_slice = jax.tree_map(lambda arr: arr[:T], obs_buf) - pi_buf_slice = pi_buf[:T] # shape [T, B, num_actions] - - # flatten [T, B] -> [T*B] - obs_dict_flat = jax.tree_map( - lambda arr: arr.reshape((T * B,) + arr.shape[2:]), - obs_dict_slice - ) - pi_buf_flat = pi_buf_slice.reshape((T * B, pi_buf.shape[-1])) - returns_buf_flat = returns_buf.reshape((T * B,)) - - # C) Train in mini-batches - - for epoch in range(config.episodes_per_iteration): - total_samples = T * B - idx = 0 - batch_size = config.batch_size - losses = [] - - while idx < total_samples: - slice_end = idx + batch_size - obs_dict_batch = jax.tree_map(lambda x: x[idx:slice_end], obs_dict_flat) - pi_batch = pi_buf_flat[idx:slice_end] - ret_batch = returns_buf_flat[idx:slice_end] - idx = slice_end - - # convert to model input - inp = obs_to_model_input(obs_dict_batch, config) - - # single step of gradient descent - 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)) - - if losses: - arr_loss = jnp.array([l[0] for l in losses]) - arr_pl = jnp.array([l[1] for l in losses]) - arr_vl = jnp.array([l[2] for l in losses]) - arr_rl = jnp.array([l[3] for l in losses]) - - mean_loss = float(arr_loss.mean()) - mean_pl = float(arr_pl.mean()) - mean_vl = float(arr_vl.mean()) - mean_rl = float(arr_rl.mean()) - else: - mean_loss = 0 - mean_pl = 0 - mean_vl = 0 - mean_rl = 0 - - wandb.log({ - "iteration": iteration, - "train/total_loss": mean_loss, - "train/policy_loss": mean_pl, - "train/value_loss": mean_vl, - "train/regularization_loss": mean_rl - }, step=iteration) - - # Evaluate raw policy - 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) - - # log action counts - wandb.log({ - "action_counts/action_0": float(eval_stats.action_0) / config.num_envs, - "action_counts/action_1": float(eval_stats.action_1) / config.num_envs, - "action_counts/action_2": float(eval_stats.action_2) / config.num_envs, - "action_counts/action_3": float(eval_stats.action_3) / config.num_envs, - "action_counts/action_4": float(eval_stats.action_4) / config.num_envs, - "action_counts/action_5": float(eval_stats.action_5) / config.num_envs, - "action_counts/action_6": float(eval_stats.action_6) / config.num_envs, - "action_counts/action_7": float(eval_stats.action_7) / config.num_envs, - "action_counts/action_8": float(eval_stats.action_8) / config.num_envs, - }, step=iteration) - - # checkpoint - if (iteration+1) % config.checkpoint_interval == 0: - ckpt = { - "iteration": iteration, - "model_params": train_state.params, - "train_config": config, - "env_config": env_params, - } - helpers.save_pkl_object(ckpt, f"checkpoints/{config.name}.pkl") - - run.finish() - return train_state - - -if __name__ == "__main__": - cfg = TrainConfig() - final_model = train_alphazero(cfg) - print("Done training!") From f87966ca9b073cb08fc1f9dbda7d6ee5a5b56a22 Mon Sep 17 00:00:00 2001 From: Chameleon01 Date: Thu, 27 Nov 2025 19:54:51 +0100 Subject: [PATCH 3/5] readme update --- README.md | 2 +- requirements.txt | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 64e93fa..9776f5a 100644 --- a/README.md +++ b/README.md @@ -161,7 +161,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/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 From 705ce6cb50869664f85fe8b57edb909503096c97 Mon Sep 17 00:00:00 2001 From: Chameleon01 Date: Thu, 27 Nov 2025 20:00:48 +0100 Subject: [PATCH 4/5] update readme --- README.md | 89 ++++++++++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 88 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 9776f5a..84b362c 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,6 +99,92 @@ wandb agent $SWEEP_ID & wait ``` +## 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 ``` From 2f54de3eeaf83c2be544c7a024d0e72ff014f75b Mon Sep 17 00:00:00 2001 From: Chameleon01 Date: Thu, 27 Nov 2025 20:06:22 +0100 Subject: [PATCH 5/5] warning added for monte carlo on readme --- README.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/README.md b/README.md index 84b362c..bee83bb 100644 --- a/README.md +++ b/README.md @@ -198,6 +198,8 @@ DATASET_PATH=/path/to/dataset DATASET_SIZE= python eval.py -r | `--debug` | Print detailed action comparisons for first 20 steps | | `-sim ` | 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