Spaces:
Running on CPU Upgrade
Running on CPU Upgrade
File size: 2,542 Bytes
2f5db5e fe3a41d 2f5db5e 6deaccc 2f5db5e 6deaccc 2f5db5e daba1b9 6deaccc 2f5db5e 6deaccc 2f5db5e | 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 | """Random baseline agent for the stellarator design environment."""
from __future__ import annotations
import random
import sys
from fusion_lab.models import StellaratorAction
from server.environment import StellaratorEnvironment
PARAMETERS = [
"aspect_ratio",
"elongation",
"rotational_transform",
"triangularity_scale",
]
DIRECTIONS = ["increase", "decrease"]
MAGNITUDES = ["small", "medium", "large"]
def random_episode(
env: StellaratorEnvironment, seed: int | None = None
) -> tuple[float, list[dict[str, object]]]:
rng = random.Random(seed)
obs = env.reset(seed=seed)
total_reward = 0.0
trace: list[dict[str, object]] = [
{
"step": 0,
"score": obs.p1_score,
"evaluation_fidelity": obs.evaluation_fidelity,
"constraints_satisfied": obs.constraints_satisfied,
}
]
while not obs.done:
if obs.budget_remaining <= 1:
action = StellaratorAction(intent="submit")
else:
action = StellaratorAction(
intent="run",
parameter=rng.choice(PARAMETERS),
direction=rng.choice(DIRECTIONS),
magnitude=rng.choice(MAGNITUDES),
)
obs = env.step(action)
total_reward += obs.reward or 0.0
trace.append(
{
"step": len(trace),
"action": action.intent,
"score": obs.p1_score,
"evaluation_fidelity": obs.evaluation_fidelity,
"constraints_satisfied": obs.constraints_satisfied,
"evaluation_failed": obs.evaluation_failed,
"reward": obs.reward,
}
)
return total_reward, trace
def main(n_episodes: int = 20) -> None:
env = StellaratorEnvironment()
rewards: list[float] = []
for i in range(n_episodes):
total_reward, trace = random_episode(env, seed=i)
final = trace[-1]
rewards.append(total_reward)
print(
f"Episode {i:3d}: steps={len(trace) - 1} "
f"final_score={final['score']:.6f} fidelity={final['evaluation_fidelity']} "
f"constraints={'yes' if final['constraints_satisfied'] else 'no'} "
f"reward={total_reward:+.4f}"
)
mean_reward = sum(rewards) / len(rewards)
print(f"\nRandom baseline ({n_episodes} episodes): mean_reward={mean_reward:+.4f}")
if __name__ == "__main__":
n = int(sys.argv[1]) if len(sys.argv) > 1 else 20
main(n)
|