donkeycar-rl-autoresearch/agent/experiments/exp28_gentrack_finetune.py

285 lines
10 KiB
Python

"""
Exp 28: Fine-tune exp26 best_model on generated track with throttle variation.
What changed from exp26:
- Warm start: exp26/best_model (best generated road model, 300k steps)
- Track: donkey-generated-track-v0 (shadows, trees) instead of generated road
- N_THROTTLE=3 (bins [0.0, 0.5, 1.0] -> clamped to [0.2, 0.5, 1.0])
exp26 used N_THROTTLE=1 (fixed throttle only). Adding throttle variation
forces the model to learn to slow into corners — critical for mini-monaco.
- Low LR=0.00005 to preserve driving skill while adapting to new visuals
- 50K steps only — just enough to adapt without forgetting road geometry
- Checkpoint every 5K, eval on generated track after each checkpoint
- After training: eval best_model on mini-monaco (zero-shot generalization test)
Goal: can adding visual diversity (shadows/trees) + throttle variation improve
generalization to mini-monaco without catastrophic forgetting?
"""
import os
import sys
import time
from datetime import datetime
sys.path.insert(0, '/home/paulh/projects/donkeycar-rl-autoresearch/agent')
_SAVE_DIR = '/home/paulh/projects/donkeycar-rl-autoresearch/agent/models/exp28-gentrack-finetune'
_PIDFILE = os.path.join(_SAVE_DIR, 'current.pid')
_WARM_MODEL = '/home/paulh/projects/donkeycar-rl-autoresearch/agent/models/exp26-warmstart/best_model.zip'
os.makedirs(_SAVE_DIR, exist_ok=True)
if os.path.exists(_PIDFILE):
try:
_old = int(open(_PIDFILE).read().strip())
if _old != os.getpid():
import signal
os.kill(_old, 0)
print(f'[exp28] Another instance already running (PID {_old}). Exiting.', flush=True)
sys.exit(1)
except (OSError, ValueError):
pass
import gymnasium as gym
import numpy as np
from stable_baselines3 import PPO
from stable_baselines3.common.vec_env import DummyVecEnv, VecTransposeImage
from discretize_action import DiscretizedActionWrapper
from donkeycar_sb3_runner import ThrottleClampWrapper
from multitrack_runner import StuckTerminationWrapper
from reward_wrapper import SpeedRewardWrapper
HOST = 'localhost'
PORT = 9091
THROTTLE_MIN = 0.2
LR = 0.00005
TOTAL_STEPS = 50_000
CHECKPOINT_EVERY = 5_000
SCENE_RELOAD_WAIT = 5.0
TRAIN_TRACK = 'donkey-generated-track-v0'
EVAL_TRACK = 'donkey-minimonaco-track-v0'
N_STEER = 7
N_THROTTLE = 1 # must match exp26 (Discrete(7)) to allow warm-start
# Same termination params as exp26
EFFICIENCY_WINDOW = 30
MIN_EFFICIENCY = 0.15
MAX_CTE = 8.0
MIN_LAP_TIME = 12.0
PROGRESS_PATIENCE = 100
MAX_STUCK_SECONDS = 5.0
MAX_EPISODE_SECONDS = 30.0
LOW_SPEED_THRESHOLD = 1.0
MAX_LOW_SPEED_SECONDS = 1.5
MAX_CTE_TERMINATION = 3.0
MAX_HIGH_CTE_SECONDS = 1.0
def log(msg):
print(f'[{datetime.now().strftime("%H:%M:%S")}] {msg}', flush=True)
def make_env(track_id, port):
def _init():
raw = gym.make(track_id, conf={'host': HOST, 'port': port})
env = ThrottleClampWrapper(raw, throttle_min=THROTTLE_MIN)
env = DiscretizedActionWrapper(env, n_steer=N_STEER, n_throttle=N_THROTTLE)
env = StuckTerminationWrapper(
env,
stuck_steps=40,
min_displacement=0.5,
max_stuck_seconds=MAX_STUCK_SECONDS,
max_episode_seconds=MAX_EPISODE_SECONDS,
low_speed_threshold=LOW_SPEED_THRESHOLD,
max_low_speed_seconds=MAX_LOW_SPEED_SECONDS,
max_cte=MAX_CTE_TERMINATION,
max_high_cte_seconds=MAX_HIGH_CTE_SECONDS,
)
env = SpeedRewardWrapper(
env,
window_size=EFFICIENCY_WINDOW,
min_efficiency=MIN_EFFICIENCY,
max_cte=MAX_CTE,
min_lap_time=MIN_LAP_TIME,
progress_patience=PROGRESS_PATIENCE,
)
return env
return _init
def connect_env(track_id=TRAIN_TRACK):
new_env = DummyVecEnv([make_env(track_id, PORT)])
new_env = VecTransposeImage(new_env)
return new_env
def reconnect_env(old_env, track_id=TRAIN_TRACK):
try:
old_env.close()
except Exception as e:
log(f' env.close() warning: {e}')
time.sleep(SCENE_RELOAD_WAIT)
return connect_env(track_id)
log('=' * 60)
log('Exp 28: gentrack fine-tune from exp26 best_model')
log(f' Sim: {HOST}:{PORT} -> {TRAIN_TRACK}')
log(f' Warm model: {_WARM_MODEL}')
log(f' Discrete: {N_STEER} steer bins, throttle fixed at {THROTTLE_MIN} (N_THROTTLE=1, matches exp26)')
log(f' LR={LR}, total={TOTAL_STEPS:,}, checkpoint every {CHECKPOINT_EVERY:,}')
log(f' After training: zero-shot eval on {EVAL_TRACK}')
log('=' * 60)
log('Connecting to sim...')
env = connect_env()
log(f' obs={env.observation_space.shape}, action={env.action_space}')
log(f'Loading warm-start model from exp26...')
model = PPO.load(_WARM_MODEL, env=env, device='cpu')
# SB3 restores lr_schedule from checkpoint; _update_learning_rate() calls
# lr_schedule(progress) each gradient step — overriding param_groups isn't enough.
# Must replace the schedule itself.
from stable_baselines3.common.utils import get_schedule_fn
model.learning_rate = LR
model.lr_schedule = get_schedule_fn(LR)
for param_group in model.policy.optimizer.param_groups:
param_group['lr'] = LR
log(f' Warm model loaded. LR={LR}')
with open(_PIDFILE, 'w') as f:
f.write(str(os.getpid()))
best_total_steps = float('-inf')
best_total_reward = float('-inf')
steps_done = 0
run_tag = datetime.now().strftime('%Y-%m-%d_%H%M%S') + '_gentrack_finetune'
log_path = os.path.join(_SAVE_DIR, f'run_{run_tag}.log')
best_model_path = os.path.join(_SAVE_DIR, 'best_model.zip')
import logging
_file_handler = logging.FileHandler(log_path)
_file_handler.setFormatter(logging.Formatter('%(message)s'))
_stream_handler = logging.StreamHandler(sys.stdout)
_stream_handler.setFormatter(logging.Formatter('%(message)s'))
file_log = logging.getLogger('exp28')
file_log.setLevel(logging.INFO)
file_log.propagate = False
file_log.addHandler(_file_handler)
file_log.addHandler(_stream_handler)
def flog(msg):
ts = datetime.now().strftime('%H:%M:%S')
file_log.info(f'[{ts}] {msg}')
flog('=' * 60)
flog(f'Exp 28 started — PID {os.getpid()}')
flog(f'Log: {log_path}')
flog(f'Warm start: exp26 best_model')
flog(f'Track: {TRAIN_TRACK} | N_STEER={N_STEER}, N_THROTTLE={N_THROTTLE}')
flog('=' * 60)
# ── Training loop ────────────────────────────────────────────────────────────
while steps_done < TOTAL_STEPS:
seg_steps = min(CHECKPOINT_EVERY, TOTAL_STEPS - steps_done)
model.learn(total_timesteps=seg_steps, reset_num_timesteps=False)
steps_done += seg_steps
ckpt = os.path.join(_SAVE_DIR, f'checkpoint_{steps_done:07d}')
model.save(ckpt)
model.save(os.path.join(_SAVE_DIR, 'model'))
flog(f'[{steps_done:,}/{TOTAL_STEPS:,}] Checkpoint saved: {ckpt}.zip')
flog(f' Reconnecting for fresh track...')
env = reconnect_env(env)
model.set_env(env)
flog(f' Connected (new track layout)')
try:
obs = env.reset()
ep_rewards = np.zeros(env.num_envs)
ep_steps = np.zeros(env.num_envs)
done_mask = np.zeros(env.num_envs, dtype=bool)
for _ in range(2000):
action, _ = model.predict(obs, deterministic=True)
obs, rewards, dones, infos = env.step(action)
for i in range(env.num_envs):
if not done_mask[i]:
ep_rewards[i] += rewards[i]
ep_steps[i] += 1
if dones[i]:
done_mask[i] = True
if done_mask.all():
break
total_steps_eval = int(ep_steps.sum())
total_reward_eval = float(ep_rewards.sum())
status = '' if ep_steps[0] >= 2000 else f'❌@{int(ep_steps[0])}'
flog(f' Eval: gentrack={total_reward_eval:.1f}r/{int(ep_steps[0])}s {status}')
if (total_steps_eval > best_total_steps
or (total_steps_eval == best_total_steps
and total_reward_eval > best_total_reward)):
best_total_steps = total_steps_eval
best_total_reward = total_reward_eval
model.save(best_model_path)
flog(f' NEW BEST: steps={best_total_steps} reward={best_total_reward:.1f}')
except Exception as e:
flog(f' Eval error: {e}')
env.close()
flog('Training complete.')
# ── Zero-shot eval on mini-monaco ────────────────────────────────────────────
flog('')
flog('=' * 60)
flog(f'ZERO-SHOT EVAL: best_model on {EVAL_TRACK}')
flog('=' * 60)
MINI_EPISODES = 5
MINI_MAX_STEPS = 3000
time.sleep(SCENE_RELOAD_WAIT)
eval_env = connect_env(track_id=EVAL_TRACK)
try:
eval_model = PPO.load(best_model_path, env=eval_env, device='cpu')
rewards_mini, steps_mini = [], []
for ep in range(1, MINI_EPISODES + 1):
obs = eval_env.reset()
total_r, steps, done = 0.0, 0, False
while not done and steps < MINI_MAX_STEPS:
action, _ = eval_model.predict(obs, deterministic=True)
obs, r, d, info = eval_env.step(action)
total_r += float(r[0])
steps += 1
done = bool(d[0])
raw_info = info[0] if isinstance(info, (list, tuple)) else info
hit = raw_info.get('hit', '?') if isinstance(raw_info, dict) else '?'
status = '✅ timeout' if steps >= MINI_MAX_STEPS else f'❌ hit={hit}@{steps}'
flog(f' ep{ep}: {total_r:.1f}r / {steps}s {status}')
rewards_mini.append(total_r)
steps_mini.append(steps)
time.sleep(0.3)
flog(f' Mean: {np.mean(steps_mini):.0f} steps / {np.mean(rewards_mini):.1f} reward')
flog(f' {"✅ GENERALIZES" if np.mean(steps_mini) > 500 else "❌ DOES NOT GENERALIZE"}')
except Exception as e:
flog(f' Mini-monaco eval error: {e}')
finally:
eval_env.close()
flog('')
flog('Exp 28 complete.')
flog(f'Log: {log_path}')