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

"""
Exp 15: Warm-start generated_track from the best mountain champion.

Goal:
- Test cross-track transfer cleanly using a single-track setup.
- Warm-start from mountain robust winner:
  agent/models/exp14-mountain-v5-finetune/best_robust_model_0036000.zip
- Train on generated_track only using the known-good Exp 13 v4 setup.

Why:
- Earlier warm-start conclusions were contaminated by bad multi-track / scene
  switching setups.
- This isolates transfer: one source model, one target track, one stable env.
"""
import sys, os, time
sys.path.insert(0, '/home/paulh/projects/donkeycar-rl-autoresearch/agent')

from donkeycar_sb3_runner import ThrottleClampWrapper
from stable_baselines3 import PPO
from stable_baselines3.common.vec_env import DummyVecEnv, VecTransposeImage
import gymnasium as gym
import numpy as np
from collections import deque
from datetime import datetime

HOST         = '10.0.0.55'
PORT         = 9091
TRACK_ID     = 'donkey-generated-track-v0'
TRACK_NAME   = 'generated_track'
THROTTLE_MIN = 0.2
SPEED_SCALE  = 0.1
LR           = 0.0004
MAX_STEPS    = 300000
EVAL_EVERY   = 5000
LAP_STOP     = 3
WARM_PATH    = '/home/paulh/projects/donkeycar-rl-autoresearch/agent/models/exp14-mountain-v5-finetune/best_robust_model_0036000.zip'
SAVE_DIR     = '/home/paulh/projects/donkeycar-rl-autoresearch/agent/models/exp15-gentrack-from-mountain'
os.makedirs(SAVE_DIR, exist_ok=True)

class V4RewardWrapper(gym.Wrapper):
    def __init__(self, env, speed_scale=0.1, window_size=60,
                 min_efficiency=0.05, max_cte=8.0):
        super().__init__(env)
        self.speed_scale    = speed_scale
        self.min_efficiency = min_efficiency
        self.max_cte        = max_cte
        self._pos_history   = deque(maxlen=window_size + 1)

    def reset(self, **kwargs):
        self._pos_history.clear()
        return self.env.reset(**kwargs)

    def step(self, action):
        result = self.env.step(action)
        if len(result) == 5:
            obs, _sim_r, terminated, truncated, info = result
            done = terminated or truncated
        else:
            obs, _sim_r, done, info = result
            terminated, truncated = done, False
        reward = self._compute_reward(done, info)
        if len(result) == 5:
            return obs, reward, terminated, truncated, info
        return obs, reward, done, info

    def _compute_reward(self, done, info):
        if done:
            return -1.0
        pos = info.get('pos', None)
        if pos is not None:
            try:
                self._pos_history.append(np.array(list(pos)[:3], dtype=np.float64))
            except (TypeError, ValueError):
                pass
        try:
            cte = float(info.get('cte', 0.0) or 0.0)
        except (TypeError, ValueError):
            cte = 0.0
        base = 1.0 - min(abs(cte) / self.max_cte, 1.0)
        efficiency = self._compute_efficiency()
        eff = max(0.0, (efficiency - self.min_efficiency) / (1.0 - self.min_efficiency))
        try:
            speed = max(0.0, float(info.get('speed', 0.0) or 0.0))
        except (TypeError, ValueError):
            speed = 0.0
        return base * eff * (1.0 + self.speed_scale * speed)

    def _compute_efficiency(self):
        if len(self._pos_history) < 3:
            return 1.0
        positions = list(self._pos_history)
        net   = np.linalg.norm(positions[-1] - positions[0])
        total = sum(np.linalg.norm(positions[i+1] - positions[i])
                    for i in range(len(positions) - 1))
        return float(net / total) if total > 1e-6 else 1.0


def log(msg):
    print(f'[{datetime.now().strftime("%H:%M:%S")}] {msg}', flush=True)


def make_env():
    def _init():
        raw = gym.make(TRACK_ID, conf={'host': HOST, 'port': PORT})
        env = ThrottleClampWrapper(raw, throttle_min=THROTTLE_MIN)
        env = V4RewardWrapper(env, speed_scale=SPEED_SCALE)
        return env
    return _init


log('='*60)
log(f'Exp 15: {TRACK_NAME} warm-start from mountain robust model')
log(f'  Host: {HOST}:{PORT}')
log(f'  Warm start: {WARM_PATH}')
log(f'  throttle_min={THROTTLE_MIN}, lr={LR}')
log(f'  Reward: v4 (Exp 13 known-good generated setup)')
log(f'  Stop: eval every {EVAL_EVERY:,} steps, stop at {LAP_STOP} laps')
log('='*60)

# Force the sim back to main menu before connecting to generated_track.
# gym.make() alone will NOT switch scenes if Unity is already running another track.
log('Switching sim to generated_track...')
_tmp = gym.make('donkey-mountain-track-v0', conf={'host': HOST, 'port': PORT})
time.sleep(2)
try:
    _tmp.unwrapped.viewer.exit_scene()
    time.sleep(0.5)
except Exception as e:
    log(f'  exit_scene warning: {e}')
_tmp.close()
time.sleep(6)
log('Sim should now be at main menu. Connecting to generated_track...')

env = VecTransposeImage(DummyVecEnv([make_env()]))
if os.path.exists(WARM_PATH):
    model = PPO.load(WARM_PATH, device='cpu')
    model.set_env(env)
    model.learning_rate = LR
    try:
        for pg in model.policy.optimizer.param_groups:
            pg['lr'] = LR
    except Exception:
        pass
    log('Loaded warm-start model and attached generated_track env')
else:
    raise FileNotFoundError(WARM_PATH)

best_reward = float('-inf')
best_laps   = 0
steps_done  = 0

while steps_done < MAX_STEPS:
    seg = min(EVAL_EVERY, MAX_STEPS - steps_done)
    model.learn(total_timesteps=seg, reset_num_timesteps=False)
    steps_done += seg

    ckpt = os.path.join(SAVE_DIR, f'checkpoint_{steps_done:07d}')
    model.save(ckpt)
    model.save(os.path.join(SAVE_DIR, 'model'))

    try:
        obs = env.reset()
        ep_r = 0.0
        ep_steps = 0
        laps = 0
        prev_lc = 0
        for _ in range(2000):
            action, _ = model.predict(obs, deterministic=True)
            obs, r, d, info = env.step(action)
            ep_r += float(r[0])
            ep_steps += 1
            try:
                lc = int((info[0] if isinstance(info, (list, tuple)) else info).get('lap_count', 0) or 0)
                if lc > prev_lc:
                    laps = lc
                    prev_lc = lc
            except Exception:
                pass
            if bool(d[0]):
                break
        status = '✅' if ep_steps >= 2000 else f'❌@{ep_steps}'
        log(f'[{steps_done:,}] reward={ep_r:.1f} steps={ep_steps} laps={laps} {status}')
        if ep_r > best_reward:
            best_reward = ep_r
            model.save(os.path.join(SAVE_DIR, 'best_model'))
            log(f'  ⭐ NEW BEST: {best_reward:.1f}')
        if laps > best_laps:
            best_laps = laps
            log(f'  🏆 BEST LAPS: {best_laps}')
        if laps >= LAP_STOP:
            log(f'  🎯 {laps} laps achieved at {steps_done:,} steps — STOPPING')
            break
    except Exception as e:
        log(f'  Eval error: {e}')

env.close()
time.sleep(3)
log(f'\nDone. best_laps={best_laps} best_reward={best_reward:.1f}')
log(f'Best model: {SAVE_DIR}/best_model.zip')
log('=== Exp 15 COMPLETE ===')