mirror of
https://github.com/velocitatem/PHANTOM.git
synced 2026-05-31 16:43:36 +00:00
337 lines
14 KiB
Python
337 lines
14 KiB
Python
"""RL training for thesis pricing system with thesis-aligned metrics.
|
||
|
||
Trains pricing policies using stable-baselines3 with TensorBoard logging.
|
||
Tracks COI erosion, alpha estimation error, and economic KPIs per thesis formulation.
|
||
|
||
Usage:
|
||
python -m lab.case.thesis.train --algo ppo --alpha 0.3 --steps 100000
|
||
python -m lab.case.thesis.train --algo adaptive --sweep # run alpha sweep
|
||
tensorboard --logdir lab/case/thesis/runs
|
||
"""
|
||
from __future__ import annotations
|
||
import argparse
|
||
import json
|
||
from dataclasses import dataclass, asdict
|
||
from pathlib import Path
|
||
from typing import Dict, List, Callable, Any
|
||
import numpy as np
|
||
|
||
try:
|
||
from stable_baselines3 import PPO, SAC, A2C
|
||
from stable_baselines3.common.callbacks import BaseCallback, EvalCallback
|
||
from stable_baselines3.common.vec_env import DummyVecEnv
|
||
from stable_baselines3.common.monitor import Monitor
|
||
HAS_SB3 = True
|
||
except ImportError:
|
||
HAS_SB3 = False
|
||
|
||
try:
|
||
from torch.utils.tensorboard import SummaryWriter
|
||
HAS_TB = True
|
||
except ImportError:
|
||
HAS_TB = False
|
||
|
||
from .simplified_env import PricingEnv, EnvConfig, make_env, adaptive_policy, fixed_price_policy, random_policy
|
||
from .simplified import coi_erosion
|
||
|
||
|
||
# thesis-aligned KPIs tracked per episode
|
||
@dataclass
|
||
class EpisodeMetrics:
|
||
reward: float = 0.0
|
||
revenue: float = 0.0
|
||
profit: float = 0.0
|
||
coi_erosion: float = 0.0 # theorem 1: order statistic erosion
|
||
coi_leakage: float = 0.0 # per-step leakage penalty
|
||
alpha_error: float = 0.0 # |α - α̂|
|
||
avg_margin: float = 0.0
|
||
n_agents: int = 0
|
||
steps: int = 0
|
||
|
||
|
||
@dataclass
|
||
class ExperimentConfig:
|
||
"""Full experiment specification for reproducibility."""
|
||
algo: str = "ppo"
|
||
total_timesteps: int = 100_000
|
||
n_envs: int = 4
|
||
eval_freq: int = 5000
|
||
n_eval_episodes: int = 10
|
||
log_dir: str = "lab/case/thesis/runs"
|
||
seed: int = 42
|
||
n_products: int = 10
|
||
max_steps: int = 200
|
||
alpha_true: float = 0.2
|
||
reward_mode: str = "robust"
|
||
experiment_name: str | None = None
|
||
|
||
def __post_init__(self):
|
||
if self.experiment_name is None:
|
||
self.experiment_name = f"{self.algo}_a{self.alpha_true:.2f}_{self.reward_mode}"
|
||
|
||
|
||
# unified policy interface wrapping all baselines
|
||
class Policy:
|
||
"""Unified policy interface for baselines and trained models."""
|
||
|
||
def __init__(self, policy_fn: Callable[[np.ndarray, int], np.ndarray], name: str):
|
||
self._fn = policy_fn
|
||
self.name = name
|
||
|
||
def predict(self, obs: np.ndarray, deterministic: bool = True) -> tuple[np.ndarray, None]:
|
||
n = (len(obs) - 3) // 3
|
||
return self._fn(obs, n), None
|
||
|
||
@staticmethod
|
||
def fixed(margin: float = 0.15) -> "Policy":
|
||
return Policy(lambda obs, n: fixed_price_policy(np.ones(n), margin), f"fixed_{margin:.2f}")
|
||
|
||
@staticmethod
|
||
def adaptive(base_margin: float = 0.15) -> "Policy":
|
||
return Policy(lambda obs, n: adaptive_policy(obs, n, base_margin), f"adaptive_{base_margin:.2f}")
|
||
|
||
@staticmethod
|
||
def random() -> "Policy":
|
||
return Policy(lambda obs, n: random_policy(n), "random")
|
||
|
||
@staticmethod
|
||
def myopic(greed: float = 0.3) -> "Policy":
|
||
"""Myopic: maximize immediate margin, ignore alpha."""
|
||
def _fn(obs: np.ndarray, n: int) -> np.ndarray:
|
||
demand_norm = obs[n:2*n] if len(obs) > 2*n else np.ones(n) * 0.5
|
||
mult = 1.0 + greed * (1 + np.mean(demand_norm))
|
||
return np.ones(n, dtype=np.float32) * np.clip(mult, 0.5, 1.5)
|
||
return Policy(_fn, f"myopic_{greed:.1f}")
|
||
|
||
|
||
class MetricsCallback(BaseCallback):
|
||
"""Tracks thesis-aligned metrics during RL training."""
|
||
|
||
def __init__(self, writer: SummaryWriter | None, verbose: int = 0):
|
||
super().__init__(verbose)
|
||
self._writer = writer
|
||
self._ep = EpisodeMetrics()
|
||
self._buffer: List[EpisodeMetrics] = []
|
||
|
||
def _on_step(self) -> bool:
|
||
for info in self.locals.get('infos', []):
|
||
self._ep.steps += 1
|
||
self._ep.reward += info.get('reward', 0)
|
||
self._ep.revenue += info.get('revenue', 0)
|
||
self._ep.profit += info.get('profit', 0)
|
||
self._ep.coi_erosion += info.get('coi_erosion', 0)
|
||
self._ep.coi_leakage += info.get('coi_leakage', 0)
|
||
self._ep.alpha_error += abs(info.get('alpha_true', 0) - info.get('alpha_est', 0))
|
||
self._ep.avg_margin += info.get('avg_margin', 0)
|
||
self._ep.n_agents += info.get('n_agents', 0)
|
||
return True
|
||
|
||
def _on_rollout_end(self) -> None:
|
||
if self._ep.steps == 0 or self._writer is None:
|
||
return
|
||
s, step = self._ep.steps, self.num_timesteps
|
||
self._writer.add_scalar('economics/revenue', self._ep.revenue / s, step)
|
||
self._writer.add_scalar('economics/profit', self._ep.profit / s, step)
|
||
self._writer.add_scalar('economics/margin', self._ep.avg_margin / s, step)
|
||
self._writer.add_scalar('coi/erosion', self._ep.coi_erosion / s, step)
|
||
self._writer.add_scalar('coi/leakage', self._ep.coi_leakage / s, step)
|
||
self._writer.add_scalar('alpha/estimation_error', self._ep.alpha_error / s, step)
|
||
self._writer.add_scalar('agents/count', self._ep.n_agents / s, step)
|
||
self._buffer.append(self._ep)
|
||
self._ep = EpisodeMetrics()
|
||
|
||
|
||
def make_vec_env(cfg: ExperimentConfig, n_envs: int = 1) -> DummyVecEnv:
|
||
def _make():
|
||
env_cfg = EnvConfig(n_products=cfg.n_products, max_steps=cfg.max_steps,
|
||
alpha_true=cfg.alpha_true, reward_mode=cfg.reward_mode, seed=cfg.seed)
|
||
return Monitor(make_env(env_cfg))
|
||
return DummyVecEnv([_make for _ in range(n_envs)])
|
||
|
||
|
||
def evaluate_policy(policy: Policy | Any, cfg: ExperimentConfig, n_episodes: int = 20) -> Dict[str, float]:
|
||
"""Evaluate policy and return thesis-aligned metrics."""
|
||
env_cfg = EnvConfig(n_products=cfg.n_products, max_steps=cfg.max_steps,
|
||
alpha_true=cfg.alpha_true, reward_mode=cfg.reward_mode, seed=cfg.seed + 999)
|
||
env = make_env(env_cfg)
|
||
metrics = []
|
||
|
||
for _ in range(n_episodes):
|
||
obs, _ = env.reset()
|
||
ep = EpisodeMetrics()
|
||
done = False
|
||
while not done:
|
||
action, _ = policy.predict(obs, deterministic=True)
|
||
obs, reward, term, trunc, info = env.step(action)
|
||
done = term or trunc
|
||
ep.reward += reward
|
||
ep.revenue += info.get('revenue', 0)
|
||
ep.profit += info.get('profit', 0)
|
||
ep.coi_erosion += info.get('coi_erosion', 0)
|
||
ep.coi_leakage += info.get('coi_leakage', 0)
|
||
ep.alpha_error += abs(info['alpha_true'] - info['alpha_est'])
|
||
ep.avg_margin += info.get('avg_margin', 0)
|
||
ep.steps += 1
|
||
metrics.append(ep)
|
||
|
||
n = len(metrics)
|
||
return {
|
||
'reward_mean': np.mean([m.reward for m in metrics]),
|
||
'reward_std': np.std([m.reward for m in metrics]),
|
||
'revenue_mean': np.mean([m.revenue / m.steps for m in metrics]),
|
||
'profit_mean': np.mean([m.profit / m.steps for m in metrics]),
|
||
'coi_erosion_mean': np.mean([m.coi_erosion / m.steps for m in metrics]),
|
||
'coi_leakage_mean': np.mean([m.coi_leakage / m.steps for m in metrics]),
|
||
'alpha_error_mean': np.mean([m.alpha_error / m.steps for m in metrics]),
|
||
'margin_mean': np.mean([m.avg_margin / m.steps for m in metrics]),
|
||
}
|
||
|
||
|
||
def train(cfg: ExperimentConfig) -> Dict[str, Any]:
|
||
"""Train RL agent or evaluate baseline policy."""
|
||
is_baseline = cfg.algo.lower() in ["fixed", "adaptive", "random", "myopic"]
|
||
if not HAS_SB3 and not is_baseline:
|
||
raise ImportError("stable-baselines3 required: pip install stable-baselines3[extra]")
|
||
|
||
log_path = Path(cfg.log_dir) / cfg.experiment_name
|
||
log_path.mkdir(parents=True, exist_ok=True)
|
||
with open(log_path / "config.json", "w") as f:
|
||
json.dump(asdict(cfg), f, indent=2)
|
||
|
||
writer = SummaryWriter(log_path) if HAS_TB else None
|
||
train_env = make_vec_env(cfg, cfg.n_envs)
|
||
eval_env = make_vec_env(cfg, 1)
|
||
|
||
if is_baseline:
|
||
policy_map = {"fixed": Policy.fixed(), "adaptive": Policy.adaptive(),
|
||
"random": Policy.random(), "myopic": Policy.myopic()}
|
||
policy = policy_map[cfg.algo.lower()]
|
||
run_baseline(policy, train_env, cfg.total_timesteps, writer)
|
||
final_metrics = evaluate_policy(policy, cfg)
|
||
else:
|
||
algo_cls = {"ppo": PPO, "sac": SAC, "a2c": A2C}.get(cfg.algo.lower())
|
||
if algo_cls is None:
|
||
raise ValueError(f"unknown algo: {cfg.algo}")
|
||
common = dict(verbose=1, seed=cfg.seed, tensorboard_log=str(log_path), device="auto")
|
||
if cfg.algo.lower() == "ppo":
|
||
model = PPO("MlpPolicy", train_env, learning_rate=3e-4, n_steps=2048,
|
||
batch_size=64, n_epochs=10, gamma=0.99, gae_lambda=0.95,
|
||
clip_range=0.2, ent_coef=0.01, **common)
|
||
elif cfg.algo.lower() == "sac":
|
||
model = SAC("MlpPolicy", train_env, learning_rate=3e-4, buffer_size=100_000,
|
||
batch_size=256, tau=0.005, gamma=0.99, **common)
|
||
else:
|
||
model = A2C("MlpPolicy", train_env, learning_rate=7e-4, n_steps=5, gamma=0.99, **common)
|
||
|
||
cb = MetricsCallback(writer)
|
||
eval_cb = EvalCallback(eval_env, best_model_save_path=str(log_path / "best"),
|
||
log_path=str(log_path), eval_freq=cfg.eval_freq,
|
||
n_eval_episodes=cfg.n_eval_episodes, deterministic=True)
|
||
model.learn(cfg.total_timesteps, callback=[cb, eval_cb], progress_bar=True)
|
||
model.save(log_path / "final_model")
|
||
policy = model
|
||
final_metrics = evaluate_policy(model, cfg)
|
||
|
||
if writer:
|
||
for k, v in final_metrics.items():
|
||
writer.add_scalar(f'final/{k}', v, cfg.total_timesteps)
|
||
writer.close()
|
||
|
||
train_env.close()
|
||
eval_env.close()
|
||
with open(log_path / "results.json", "w") as f:
|
||
json.dump(final_metrics, f, indent=2)
|
||
return {"path": str(log_path), "metrics": final_metrics}
|
||
|
||
|
||
def run_baseline(policy: Policy, vec_env: DummyVecEnv, total_steps: int, writer: SummaryWriter | None):
|
||
"""Run baseline policy through environment with logging."""
|
||
obs = vec_env.reset()
|
||
n_envs = vec_env.num_envs
|
||
ep_rewards = np.zeros(n_envs)
|
||
all_rewards, coi_buf, alpha_buf = [], [], []
|
||
|
||
for step in range(0, total_steps, n_envs):
|
||
actions = np.array([policy.predict(obs[i])[0] for i in range(n_envs)])
|
||
obs, rewards, dones, infos = vec_env.step(actions)
|
||
ep_rewards += rewards
|
||
for i, info in enumerate(infos):
|
||
coi_buf.append(info.get('coi_erosion', 0))
|
||
alpha_buf.append(abs(info.get('alpha_true', 0) - info.get('alpha_est', 0)))
|
||
if dones[i]:
|
||
all_rewards.append(ep_rewards[i])
|
||
ep_rewards[i] = 0
|
||
if writer and step % 1000 < n_envs and all_rewards:
|
||
writer.add_scalar('rollout/ep_rew_mean', np.mean(all_rewards[-20:]), step)
|
||
writer.add_scalar('coi/erosion', np.mean(coi_buf[-100:]), step)
|
||
writer.add_scalar('alpha/estimation_error', np.mean(alpha_buf[-100:]), step)
|
||
|
||
|
||
def run_sweep(cfg: ExperimentConfig, alphas: List[float] | None = None) -> Dict[str, Dict]:
|
||
"""Run experiment across contamination levels for scientific comparison."""
|
||
alphas = alphas or [0.0, 0.1, 0.2, 0.3, 0.4, 0.5]
|
||
results = {}
|
||
for alpha in alphas:
|
||
sweep_cfg = ExperimentConfig(**{**asdict(cfg), "alpha_true": alpha,
|
||
"experiment_name": f"{cfg.algo}_a{alpha:.2f}_{cfg.reward_mode}"})
|
||
print(f"\n=== α={alpha:.2f} ===")
|
||
out = train(sweep_cfg)
|
||
results[f"alpha_{alpha:.2f}"] = out["metrics"]
|
||
summary_path = Path(cfg.log_dir) / f"sweep_{cfg.algo}_{cfg.reward_mode}.json"
|
||
with open(summary_path, "w") as f:
|
||
json.dump(results, f, indent=2)
|
||
print(f"\nSweep results saved to {summary_path}")
|
||
return results
|
||
|
||
|
||
def compare_policies(cfg: ExperimentConfig, policies: List[str] | None = None) -> Dict[str, Dict]:
|
||
"""Compare multiple policies at same contamination level."""
|
||
policies = policies or ["fixed", "adaptive", "myopic", "random"]
|
||
results = {}
|
||
for algo in policies:
|
||
cmp_cfg = ExperimentConfig(**{**asdict(cfg), "algo": algo,
|
||
"experiment_name": f"cmp_{algo}_a{cfg.alpha_true:.2f}"})
|
||
print(f"\n=== {algo} ===")
|
||
out = train(cmp_cfg)
|
||
results[algo] = out["metrics"]
|
||
cmp_path = Path(cfg.log_dir) / f"compare_a{cfg.alpha_true:.2f}.json"
|
||
with open(cmp_path, "w") as f:
|
||
json.dump(results, f, indent=2)
|
||
print(f"\nComparison saved to {cmp_path}")
|
||
for algo, m in results.items():
|
||
print(f" {algo:12s}: reward={m['reward_mean']:.2f} coi_erosion={m['coi_erosion_mean']:.4f} "
|
||
f"alpha_err={m['alpha_error_mean']:.4f}")
|
||
return results
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="Train RL pricing policies")
|
||
parser.add_argument("--algo", default="ppo", choices=["ppo", "sac", "a2c", "fixed", "adaptive", "random", "myopic"])
|
||
parser.add_argument("--steps", type=int, default=100_000)
|
||
parser.add_argument("--alpha", type=float, default=0.2)
|
||
parser.add_argument("--reward-mode", default="robust", choices=["revenue", "profit", "robust", "coi_aware"])
|
||
parser.add_argument("--n-products", type=int, default=10)
|
||
parser.add_argument("--n-envs", type=int, default=4)
|
||
parser.add_argument("--seed", type=int, default=42)
|
||
parser.add_argument("--log-dir", default="lab/case/thesis/runs")
|
||
parser.add_argument("--sweep", action="store_true", help="run contamination sweep")
|
||
parser.add_argument("--compare", action="store_true", help="compare all baselines")
|
||
args = parser.parse_args()
|
||
|
||
cfg = ExperimentConfig(algo=args.algo, total_timesteps=args.steps, alpha_true=args.alpha,
|
||
reward_mode=args.reward_mode, n_products=args.n_products,
|
||
n_envs=args.n_envs, seed=args.seed, log_dir=args.log_dir)
|
||
|
||
if args.sweep:
|
||
run_sweep(cfg)
|
||
elif args.compare:
|
||
compare_policies(cfg)
|
||
else:
|
||
result = train(cfg)
|
||
print(f"\nTraining complete: {result['path']}")
|
||
print(f"Metrics: {json.dumps(result['metrics'], indent=2)}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|