Spaces:
Running
Running
| 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() | |