File size: 1,906 Bytes
2e739de | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | """
Run CSPO vs PPO-Lag on SafetyPointGoal1-v0 for comparison.
Reduced scale: 200K steps (paper uses 10M).
"""
import warnings
warnings.filterwarnings('ignore')
import os
os.environ['WANDB_MODE'] = 'disabled'
import omnisafe
import numpy as np
def run_experiment(algo, env_id, seed, total_steps, log_dir):
agent = omnisafe.Agent(
algo,
env_id,
custom_cfgs={
'seed': seed,
'train_cfgs': {
'total_steps': total_steps,
'vector_env_nums': 1,
'torch_threads': 4,
'device': 'cpu',
},
'algo_cfgs': {
'steps_per_epoch': 20000,
'update_iters': 10,
'batch_size': 512,
},
'logger_cfgs': {
'use_wandb': False,
'use_tensorboard': False,
'log_dir': log_dir,
'window_lens': 10,
},
},
)
agent.learn()
# Read final metrics from the logger
final_ret = agent._logger.get_stats('Metrics/EpRet')
final_cost = agent._logger.get_stats('Metrics/EpCost')
return final_ret, final_cost
# Run CSPO on PointGoal
print("Running CSPO on SafetyPointGoal1-v0...")
cspo_ret, cspo_cost = run_experiment(
'CSPO', 'SafetyPointGoal1-v0', seed=0,
total_steps=200000, log_dir='/tmp/cspo_bench/cspo'
)
print(f"CSPO - Return: {cspo_ret:.2f}, Cost: {cspo_cost:.2f}")
# Run PPO-Lag on PointGoal
print("\nRunning PPO-Lag on SafetyPointGoal1-v0...")
ppolag_ret, ppolag_cost = run_experiment(
'PPOLag', 'SafetyPointGoal1-v0', seed=0,
total_steps=200000, log_dir='/tmp/cspo_bench/ppolag'
)
print(f"PPO-Lag - Return: {ppolag_ret:.2f}, Cost: {ppolag_cost:.2f}")
print("\nComparison:")
print(f"CSPO: Return={cspo_ret:.2f}, Cost={cspo_cost:.2f}")
print(f"PPO-Lag: Return={ppolag_ret:.2f}, Cost={ppolag_cost:.2f}")
|