Spaces:
Running
Running
File size: 3,719 Bytes
ad77c79 | 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 | import json
import random
from server.openenv_wrapper import SupplyChainEnv
from server.engine import CONTAINER_SIZE, SHIPPING_COST
def generate_expert_trajectories(num_episodes=20, output_file="dataset.jsonl"):
env = SupplyChainEnv(max_turns=24)
dataset = []
print(f"Generating {num_episodes} expert episodes...")
for episode in range(num_episodes):
obs_json, info = env.reset()
done = False
while not done:
obs = json.loads(obs_json)
state_dict = env._state # Accessing internal state for heuristic planning
# Find available ports
open_ports = [p for p in state_dict.portDelays if not p["isClosed"]]
if not open_ports:
# No open ports, only possible action is to advance
action_dict = {"action": "advance"}
# Log the pair
dataset.append({
"prompt": f"System: You are an AI Logistics Agent. Output only JSON.\nObservation: {obs_json}\nAction: ",
"completion": json.dumps(action_dict)
})
obs_json, reward, done, truncated, info = env.step(json.dumps(action_dict))
continue
# Best port and factory heuristic
best_port = min(open_ports, key=lambda x: x["currentDelay"])
best_factory = min(state_dict.factoryCosts, key=lambda x: x["currentCost"])
actions_taken_this_turn = False
for region, demand in state_dict.demands.items():
if state_dict.inventory[region] < demand:
needed = demand - state_dict.inventory[region]
cost_of_goods = best_factory["currentCost"] * needed
containers = (needed + CONTAINER_SIZE - 1) // CONTAINER_SIZE
total_cost = cost_of_goods + (containers * SHIPPING_COST)
if state_dict.balance >= total_cost:
action_dict = {
"action": "order",
"factory": best_factory["id"],
"quantity": needed,
"port": best_port["id"],
"dest": region
}
# Log the order action
dataset.append({
"prompt": f"System: You are an AI Logistics Agent. Output only JSON.\nObservation: {obs_json}\nAction: ",
"completion": json.dumps(action_dict)
})
obs_json, reward, done, truncated, info = env.step(json.dumps(action_dict))
actions_taken_this_turn = True
# After fulfilling all possible demands, we must advance the turn
advance_action = {"action": "advance"}
dataset.append({
"prompt": f"System: You are an AI Logistics Agent. Output only JSON.\nObservation: {obs_json}\nAction: ",
"completion": json.dumps(advance_action)
})
obs_json, reward, done, truncated, info = env.step(json.dumps(advance_action))
# Write to JSONL
with open(output_file, "w") as f:
for entry in dataset:
f.write(json.dumps(entry) + "\n")
print(f"✅ Generated {len(dataset)} expert action pairs and saved to {output_file}")
if __name__ == "__main__":
generate_expert_trajectories()
|