File size: 4,504 Bytes
99d2ff3 | 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 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | import numpy as np
import gymnasium as gym
from gymnasium import spaces
from models import Action
from server.llm_env import LLMEnv
from stable_baselines3 import PPO
# Define the precise mapping of numerical indices to our Pydantic string actions
ACTION_MAPPING = [
"follow_prompt",
"process_data",
"routine_eval",
"seek_feedback",
"minor_hallucination",
"major_hallucination",
"cascade_hallucination",
"optimize_context",
"lower_temperature"
]
class LLMGymWrapper(gym.Env):
"""
Wraps the OpenEnv LLMEnv into a standard Gymnasium Environment
compatible with stable-baselines3 algorithms.
"""
def __init__(self, task="medium", max_days=180):
super(LLMGymWrapper, self).__init__()
self.env = LLMEnv(task=task, max_days=max_days)
# 9 Discrete actions
self.action_space = spaces.Discrete(len(ACTION_MAPPING))
# 8 Observation variables (day, alignment, hallucination, user_trust, entropy_level, compute, moderation, filter_risk)
# All bounded within [0, infinity] for safety, although most are 0-100.
self.observation_space = spaces.Box(
low=np.zeros(8, dtype=np.float32),
high=np.full(8, np.inf, dtype=np.float32),
dtype=np.float32
)
def _get_obs_array(self, obs) -> np.ndarray:
return np.array([
obs.day,
obs.alignment,
obs.hallucination,
obs.user_trust,
obs.entropy_level,
obs.compute,
obs.moderation,
obs.filter_risk
], dtype=np.float32)
def reset(self, seed=None, options=None):
super().reset(seed=seed)
if seed is not None:
self.env.seed = seed
self.env.rng.seed(seed)
obs = self.env.reset()
return self._get_obs_array(obs), {}
def step(self, action_idx: int):
# Decode action index to the specific Pydantic string
action_type = ACTION_MAPPING[action_idx]
action = Action(action_type=action_type)
obs_obj, reward, done, info = self.env.step(action)
obs = self._get_obs_array(obs_obj)
# In Gymnasium, done is split into terminated and truncated.
# OpenEnv 'done' handles both, we'll map done to terminated for simplicity.
terminated = done
truncated = False
return obs, reward, terminated, truncated, info
def train_and_evaluate():
print("Initializing LLM Environment Gymnasium Wrapper...")
env = LLMGymWrapper(task="medium")
# Initialize PPO
print("Initializing PPO Agent...")
model = PPO("MlpPolicy", env, verbose=1)
# Train the agent
print("Training Agent for 20,000 steps to learn alignment...")
model.learn(total_timesteps=20000)
# Save the model
model.save("ppo_llm_aligned")
print("\n--- Training Complete ---")
print("Evaluating the trained model on an episode to verify hallucination-avoidance.\n")
# Evaluation
obs, info = env.reset()
done = False
total_reward = 0.0
action_counts = {a: 0 for a in ACTION_MAPPING}
while not done:
action, _states = model.predict(obs, deterministic=True)
obs, reward, terminated, truncated, info = env.step(action.item())
done = terminated or truncated
total_reward += reward
action_str = ACTION_MAPPING[action.item()]
action_counts[action_str] += 1
print(f"Evaluation Complete! Episode ended at day {obs[0]}")
print(f"Cumulative Reward: {total_reward:.2f}")
print("\nAction Selection Frequency:")
for a, count in action_counts.items():
if count > 0:
print(f" - {a}: {count} times")
print("\nFinal State Values:")
print(f" - Alignment: {obs[1]:.1f}%")
print(f" - Hallucination Level: {obs[2]:.1f}%")
print(f" - Compute Level: {obs[5]:.1f}%")
print(f" - Filter Risk: {obs[7]:.1f}%")
# Analyze alignment
bad_actions = ["minor_hallucination", "major_hallucination", "cascade_hallucination"]
hallucination_count = sum(action_counts[a] for a in bad_actions)
if hallucination_count == 0:
print("\nSUCCESS: The agent learned to perfectly avoid hallucinating actions and align with the user!")
else:
print(f"\nNOTE: The agent still hallucinated {hallucination_count} times.")
if __name__ == "__main__":
train_and_evaluate()
|