tasks/hard.py

from __future__ import annotations

from typing import Any, Dict

from env.models import EnvironmentConfig


TASK_NAME = "hard"
TASK_DIFFICULTY = 3
DEFAULT_SEED = 303
DESCRIPTION = "High-pressure dispatch benchmark with strict energy limits, priority SLAs, and tight step budget."
OBJECTIVE = "Sustain high on-time completion for priority-heavy workloads under battery, traffic, and dynamic obstacle pressure."
SUCCESS_CONDITION: Dict[str, Any] = {
	"completion_rate_min": 0.90,
	"high_priority_on_time_rate_min": 0.88,
	"battery_depletion": False,
	"max_steps": 165,
	"invalid_action_rate_max": 0.04,
}

TASK_PARAMETERS: Dict[str, Any] = {
	"grid_size": {"width": 16, "height": 16},
	"orders": {"min": 6, "max": 8},
	"variation": {
		"seeded_layout": True,
		"order_mix": "priority_heavy",
		"traffic_profile": "peak_hour",
		"battery_profile": "tight_capacity",
		"step_budget": "tight",
	},
	"constraints": {
		"obstacles": True,
		"dynamic_obstacles": True,
		"traffic_zones": True,
		"order_priority": True,
		"multiple_delivery_locations": True,
		"battery_constraint": True,
		"recharge": True,
		"time_penalty": True,
	},
}


def get_config(seed: int | None = None) -> EnvironmentConfig:
	effective_seed = DEFAULT_SEED if seed is None else seed
	order_count = 6 + (effective_seed % 3)
	delivery_locations = 2 + (effective_seed % 2)
	dynamic_ratio = 0.50 + (0.05 * (effective_seed % 3))

	return EnvironmentConfig(
		width=16,
		height=16,
		max_steps=165,
		max_orders=order_count,
		delivery_locations_per_order=delivery_locations,
		priority_high_ratio=0.75,
		obstacle_density=0.20,
		dynamic_obstacles_enabled=True,
		dynamic_obstacle_ratio=min(0.65, dynamic_ratio),
		traffic_density=0.22,
		traffic_extra_cost=5,
		battery_enabled=True,
		battery_capacity=55,
		battery_recharge_rate=8,
		seed=effective_seed,
	)


def get_definition() -> Dict[str, Any]:
	return {
		"name": TASK_NAME,
		"difficulty": TASK_DIFFICULTY,
		"deterministic": True,
		"default_seed": DEFAULT_SEED,
		"description": DESCRIPTION,
		"objective": OBJECTIVE,
		"success_condition": SUCCESS_CONDITION,
		"parameters": TASK_PARAMETERS,
	}