there was a bug...

.gitignore

@@ -215,5 +215,4 @@ __marimo__/
 # Streamlit
 .streamlit/secrets.toml
 
-charts/
-!charts/gantt_fifo_policy.png
+output/

README.md

@@ -53,7 +53,7 @@ The project follows a client-server architecture using the OpenEnv framework:
 
 **Models** (`src/jssp_openenv/models.py`):
 - `JSSPAction`: Represents scheduling actions (list of job IDs to schedule)
-- `JSSPObservation`: Contains the current state (machines, ready operations, progress)
+- `JSSPObservation`: Contains the current state (machine status, job status, remaining operations)
 
 **Environment** (`src/jssp_openenv/server/jssp_environment.py`):
 - `JSSPEnvironment`: The core simulation environment that:
@@ -144,6 +144,20 @@ Here is an example:
 
 ![FIFO Policy Gantt Chart](assets/gantt_fifo_policy.png)
 
+## Current Results
+
+Results as of Nov. 7, 2024 on FT06 problem instance. *Note: Non-scientific results, only ran 1 episode per policy.*
+
+| Policy | Makespan |
+|--------|----------|
+| **Optimal solution** | **55** |
+| `openai/gpt-oss-20b:groq` | 61 |
+| FIFO | 68 |
+| `openai/gpt-oss-120b:cerebras` | 69 |
+| `Qwen/Qwen3-32B:groq` | 69 |
+| Max-Min | 77 |
+
+
 ## Run with docker
 
 Build the Docker image:

pyproject.toml

@@ -26,11 +26,9 @@ dev = [
     "ty",
 ]
 
-[tool.setuptools]
-package-dir = {"" = "src"}
-
-[tool.setuptools.packages.find]
-where = ["src"]
+[tool.mypy]
+disable_error_code = ["import-untyped"]
+ignore_missing_imports = true
 
 [tool.ruff]
 exclude = [".git", ".ruff_cache", ".venv"]
@@ -38,4 +36,10 @@ line-length = 119
 # Ignored rules:
 # "E501" -> line length violation
 lint.ignore = ["E501"]
-lint.select = ["E", "F", "I", "W"]
+lint.select = ["E", "F", "I", "W"]
+
+[tool.setuptools]
+package-dir = {"" = "src"}
+
+[tool.setuptools.packages.find]
+where = ["src"]

run.py

@@ -11,8 +11,8 @@ from jssp_openenv.solver import solve_jssp
 
 SERVER_URL = "http://localhost:8000"
 MAX_STEPS = 1000  # Maximum number of steps per instance
-CHART_DIR = "charts"
-os.makedirs(CHART_DIR, exist_ok=True)
+OUTPUT_DIR = "output"
+os.makedirs(OUTPUT_DIR, exist_ok=True)
 
 cli = typer.Typer()
 
@@ -68,7 +68,7 @@ def solve(
 
     print(f"Solved in {makespan} steps")
 
-    filepath = os.path.join(CHART_DIR, filename)
+    filepath = os.path.join(OUTPUT_DIR, filename)
     gantt_chart(scheduled_events, title=title, makespan=makespan, save_to=filepath)
     print(f"Saved Gantt chart to {filepath}")
 

src/jssp_openenv/client.py

@@ -1,6 +1,6 @@
 from openenv_core import HTTPEnvClient, StepResult
 
-from .models import JSSPAction, JSSPObservation, MachineObservation, ReadyOperationObservation
+from .models import JobObservation, JSSPAction, JSSPObservation, MachineObservation
 
 
 class JSSPEnvClient(HTTPEnvClient[JSSPAction, JSSPObservation]):
@@ -12,9 +12,7 @@ class JSSPEnvClient(HTTPEnvClient[JSSPAction, JSSPObservation]):
         return StepResult[JSSPObservation](
             observation=JSSPObservation(
                 machines=[MachineObservation(**machine) for machine in obs_data.pop("machines")],
-                ready_operations=[
-                    ReadyOperationObservation(**operation) for operation in obs_data.pop("ready_operations")
-                ],
+                jobs=[JobObservation(**job) for job in obs_data.pop("jobs")],
                 **obs_data,
             ),
             reward=payload.get("reward"),

src/jssp_openenv/models.py

@@ -32,11 +32,12 @@ class MachineObservation:
 
 
 @dataclass
-class ReadyOperationObservation:
+class JobObservation:
+    """Observation of a given Job in the JSSP environment."""
+
     job_id: int
-    machine_id: int
-    duration: int
-    remaining_ops: int
+    operations: JobT  # remaining operations to be scheduled
+    busy_until: Optional[int]  # time until the current operation is complete (or none if available)
 
 
 @dataclass(kw_only=True)
@@ -47,17 +48,30 @@ class JSSPObservation(Observation):
 
     step_count: int
     machines: list[MachineObservation]
-    ready_operations: list[ReadyOperationObservation]
-    completed_jobs: int
-    total_jobs: int
+    jobs: list[JobObservation]
+
+    def available_machines(self) -> list[MachineObservation]:
+        """Get available machines from observation."""
+        return [m for m in self.machines if m.busy_until is None or m.busy_until <= self.step_count]
+
+    def available_jobs(self) -> list[JobObservation]:
+        """Get available jobs from observation."""
+        available_machine_ids = [m.machine_id for m in self.available_machines()]
+        return [
+            job
+            for job in self.jobs
+            if (job.busy_until is None or job.busy_until <= self.step_count)
+            and len(job.operations) > 0
+            and job.operations[0][0] in available_machine_ids
+        ]
 
 
 def parse_job_ids(job_ids: str) -> list[int]:
     """Parse job_ids from string (error out if cannot be parsed)."""
     try:
         return [int(job_id) for job_id in job_ids.split(",") if job_id.strip()]
-    except ValueError:
-        raise ValueError(f"Invalid job_ids: {job_ids}")
+    except ValueError as e:
+        raise ValueError(f"Invalid job_ids: {job_ids}") from e
 
 
 @dataclass

src/jssp_openenv/policy.py

@@ -3,7 +3,7 @@ from abc import ABC, abstractmethod
 
 from openai import OpenAI
 
-from .models import JSSPAction, JSSPObservation, MachineObservation, ReadyOperationObservation
+from .models import JobObservation, JSSPAction, JSSPObservation, MachineObservation
 
 
 class JSSPEnvPolicy(ABC):
@@ -15,30 +15,24 @@ class JSSPEnvPolicy(ABC):
 class JSSPFifoPolicy(JSSPEnvPolicy):
     def act(self, observation: JSSPObservation) -> JSSPAction:
         """
-        FIFO scheduling: schedule ready operations in order of job_id.
+        FIFO scheduling: schedule available jobs in order of job_id.
 
-        This policy schedules operations in FIFO order (by job_id), respecting
-        machine availability. It only schedules operations for machines that are
-        currently available (not busy).
+        This policy schedules jobs in FIFO order (by job_id), respecting machine availability.
+        It only schedules jobs for machines that are currently available (not busy).
         """
-        # Create a lookup for machine availability
-        machine_available = {m.machine_id: m.busy_until is None for m in observation.machines}
-
-        # Filter to only ready operations with available machines
-        available_ops = [op for op in observation.ready_operations if machine_available.get(op.machine_id, False)]
-
-        # Sort by job_id (FIFO: first job_id first)
-        available_ops.sort(key=lambda op: op.job_id)
+        # Filter to only available jobs with available machines + sort by job_id
+        sorted_jobs = sorted(observation.available_jobs(), key=lambda job: job.job_id)
 
         # Track which machines we've already scheduled to avoid conflicts
         scheduled_machines = set()
         scheduled_job_ids = []
 
         # Schedule jobs in FIFO order, but skip if machine is already taken
-        for op in available_ops:
-            if op.machine_id not in scheduled_machines:
-                scheduled_job_ids.append(op.job_id)
-                scheduled_machines.add(op.machine_id)
+        for job in sorted_jobs:
+            machine_id = job.operations[0][0]
+            if machine_id not in scheduled_machines:
+                scheduled_job_ids.append(job.job_id)
+                scheduled_machines.add(machine_id)
 
         return JSSPAction(job_ids=scheduled_job_ids)
 
@@ -48,18 +42,19 @@ class JSSPMaxMinPolicy(JSSPEnvPolicy):
         """
         Max-Min scheduling: schedule the operation with the longest duration first.
         """
-        # Sort operations by duration (max-min)
-        ops = sorted(observation.ready_operations, key=lambda op: op.duration, reverse=True)
+        # Sort available jobs by duration (max-min)
+        sorted_jobs = sorted(observation.available_jobs(), key=lambda job: job.operations[0][1], reverse=True)
 
         # Track which machines we've already scheduled to avoid conflicts
         scheduled_machines = set()
         scheduled_job_ids = []
 
         # Schedule jobs in max-min order, but skip if machine is already taken
-        for op in ops:
-            if op.machine_id not in scheduled_machines:
-                scheduled_job_ids.append(op.job_id)
-                scheduled_machines.add(op.machine_id)
+        for job in sorted_jobs:
+            machine_id = job.operations[0][0]
+            if machine_id not in scheduled_machines:
+                scheduled_job_ids.append(job.job_id)
+                scheduled_machines.add(machine_id)
 
         return JSSPAction(job_ids=scheduled_job_ids)
 
@@ -69,46 +64,33 @@ You are solving a Job Shop Scheduling Problem (JSSP). Your goal is to minimize t
 
 You must optimize for minimal makespan while respecting all constraints. Each job consists of multiple operations that must be completed in sequence, and each operation requires a specific machine for a given duration.
 
----
-
-### 🕒 Current State
-**Step:** {step_count} | **Completed:** {completed_jobs}/{total_jobs}
+**Current step:** {step_count}
 
----
-
-### ⚙️ Machine Status
+**Machine Status:**
 {machines_status}
 
 You must check machine availability before scheduling. Machines that are busy cannot start new operations until they finish their current task.
 
----
-
-### ✅ Ready to Schedule (NOW)
-{ready_operations_list}
-
-Each entry shows: **machine**, **duration**, and **remaining ops**.
-You can only schedule operations that are ready at this step. These are operations whose previous steps in the job sequence have been completed.
+**Jobs:**
+{jobs_list}
 
----
-
-### 🎯 Rules You Must Follow
-1. You must schedule only **ready** operations. Do not attempt to schedule operations that are not ready.
+**Rules You Must Follow:**
+1. You must schedule only **available** jobs. Do not attempt to schedule jobs that are not available.
 2. Each machine can run **one job at a time**. You cannot schedule multiple jobs on the same machine simultaneously.
 3. You must not schedule jobs on **busy** machines (`busy_until > current step`). Check machine availability before scheduling.
 4. You may **schedule multiple** jobs on different machines in the same step, or you may choose to wait if no good scheduling opportunity exists.
 
----
-
-### 🧩 Available Actions
-{legal_actions}
+**Legal actions:**
+{legal_job_ids}
 
-These are the valid job IDs you can schedule at this step. You must choose from this list.
+These are the valid job IDs you can schedule at this step. You must choose a subset from this list or choose to wait.
 
-**Answer format:**
+**Action format:**
 - To schedule jobs: `"0,2"` or `"1"` (comma-separated job IDs)
 - To wait: `""` (empty string)
+Select the best subset of jobs to schedule to minimize the total makespan once all jobs are completed.
 
-Respond only with the action format specified above.
+Response only with the action format specified above, and nothing else.
 """
 
 
@@ -132,35 +114,27 @@ class JSSPLLMPolicy(JSSPEnvPolicy):
         """
         LLM scheduling: use an LLM to schedule the operations.
 
-        Determines legal actions (ready operations with available machines),
-        formats a prompt, calls the LLM, and parses the response to return
-        a scheduling action.
+        Process:
+        - Determine legal job IDs (available jobs on available machines)
+        - Format a prompt
+        - Call the LLM
+        - Parse the response to return a scheduling action
         """
-        # Determine machine availability
-        machine_available = {
-            m.machine_id: m.busy_until is None or m.busy_until <= observation.step_count for m in observation.machines
-        }
-
-        # Filter ready operations to only include those with available machines
-        legal_job_ids = [
-            op.job_id for op in observation.ready_operations if machine_available.get(op.machine_id, False)
-        ]
+        available_jobs = observation.available_jobs()
 
         # If no legal actions, return empty action (wait)
-        if not legal_job_ids:
+        if not available_jobs:
             return JSSPAction(job_ids=[])
 
         # Format prompt
         machines_status = self._format_machines_status(observation.machines, observation.step_count)
-        ready_operations_list = self._format_ready_operations(observation.ready_operations)
+        jobs_list = self._format_jobs(observation.jobs)
 
         prompt = PROMPT_TEMPLATE.format(
             step_count=observation.step_count,
-            completed_jobs=observation.completed_jobs,
-            total_jobs=observation.total_jobs,
             machines_status=machines_status,
-            ready_operations_list=ready_operations_list,
-            legal_actions=legal_job_ids,
+            jobs_list=jobs_list,
+            legal_job_ids=[job.job_id for job in available_jobs],
         )
         print(f"Step {observation.step_count}")
 
@@ -170,9 +144,7 @@ class JSSPLLMPolicy(JSSPEnvPolicy):
                 model=self.model_id, messages=[{"role": "user", "content": prompt}], temperature=0.0
             )
             llm_output = response.choices[0].message.content or ""
-            print(f"LLM Output: {llm_output}")
-            job_ids = self._parse_action(llm_output, legal_job_ids)
-            print(f"Job IDs: {job_ids}")
+            job_ids = self._parse_action(llm_output, available_jobs)
 
             # Ensure we don't schedule multiple jobs on the same machine
             # Track which machines we've already scheduled to avoid conflicts
@@ -180,10 +152,10 @@ class JSSPLLMPolicy(JSSPEnvPolicy):
             filtered_job_ids = []
             for job_id in job_ids:
                 # Find the operation for this job
-                op = next((op for op in observation.ready_operations if op.job_id == job_id), None)
-                if op is not None and op.machine_id not in scheduled_machines:
+                op = next((op for op in available_jobs if op.job_id == job_id), None)
+                if op is not None and op.operations[0][0] not in scheduled_machines:
                     filtered_job_ids.append(job_id)
-                    scheduled_machines.add(op.machine_id)
+                    scheduled_machines.add(op.operations[0][0])
 
             return JSSPAction(job_ids=filtered_job_ids)
 
@@ -207,18 +179,23 @@ class JSSPLLMPolicy(JSSPEnvPolicy):
         return "\n".join(lines) if lines else "  (No machines)"
 
     @staticmethod
-    def _format_ready_operations(ready_operations: list[ReadyOperationObservation]) -> str:
-        """Format ready operations for prompt."""
+    def _format_jobs(jobs: list[JobObservation]) -> str:
+        """Format jobs for prompt."""
         lines = []
-        for op in ready_operations:
-            lines.append(
-                f"  Job {op.job_id}: Machine {op.machine_id}, Duration {op.duration} min, {op.remaining_ops} ops remaining"
-            )
-        return "\n".join(lines) if lines else "  (No ready operations)"
+        for job in jobs:
+            available = job.busy_until is None
+            operations = ", ".join(f"(Machine {op[0]}, {op[1]} min)" for op in job.operations)
+            if available:
+                lines.append(f"  Job {job.job_id}: Available. Remaining operations: {operations}")
+            else:
+                lines.append(f"  Job {job.job_id}: Busy until t={job.busy_until}. Remaining operations: {operations}")
+        return "\n".join(lines) if lines else "  (No jobs)"
 
     @staticmethod
-    def _parse_action(text: str, legal_job_ids: list[int]) -> list[int]:
+    def _parse_action(text: str, available_jobs: list[JobObservation]) -> list[int]:
         """Parse comma-separated job IDs from model output."""
+        legal_job_ids = [job.job_id for job in available_jobs]
+
         # First, we remove the reasoning section
         text = text.split("<think>")[-1].split("</think>")[-1].strip()
 

src/jssp_openenv/server/jssp_environment.py

@@ -5,7 +5,7 @@ from typing import Optional
 import simpy
 from openenv_core.env_server import Environment
 
-from ..models import JobT, JSSPAction, JSSPObservation, MachineObservation, ReadyOperationObservation
+from ..models import JobObservation, JobT, JSSPAction, JSSPObservation, MachineObservation
 
 # Example of JSSP initial jobs
 # Each tuple is a (machine_index, processing_time)
@@ -50,35 +50,33 @@ class JSSPEnvironment(Environment):
 
         return self.state
 
-    def _get_ready_operations(self) -> list[ReadyOperationObservation]:
-        """Get all operations that are ready to be scheduled now."""
-        ready = []
+    def _get_jobs(self) -> list[JobObservation]:
+        """Get all jobs with their status and remaining operations."""
+        jobs: list[JobObservation] = []
         for job_id in range(len(self.jobs)):
-            # Skip if job is complete
-            if self.job_progress[job_id] >= len(self.jobs[job_id]):
-                continue
+            # @dataclass
+            # class JobObservation:
+            #     """Observation of a given Job in the JSSP environment."""
 
-            # Get next operation for this job
-            machine_id, duration = self.jobs[job_id][self.job_progress[job_id]]
+            #     job_id: int
+            #     operations: JobT  # remaining operations to be scheduled (not counting the current one)
+            #     busy_until: Optional[int]  # time until the current operation is complete (or none if available)
+            job_operations = self.jobs[job_id]
+            job_progress = self.job_progress[job_id]
+            job_remaining_operations = job_operations[job_progress:]
 
-            # Check if machine is available
-            busy_until = self.machine_busy_until[machine_id]
-            if busy_until is None or busy_until <= self.env.now:
-                remaining = len(self.jobs[job_id]) - self.job_progress[job_id]
-                ready.append(
-                    ReadyOperationObservation(
-                        job_id=job_id,
-                        machine_id=machine_id,
-                        duration=duration,
-                        remaining_ops=remaining,
-                    )
-                )
-
-        return ready
+            job_busy_until = None
+            for current_job, busy_until in zip(self.machine_current_job, self.machine_busy_until):
+                if current_job is not None and current_job == job_id:
+                    job_busy_until = busy_until
+
+            jobs.append(JobObservation(job_id=job_id, operations=job_remaining_operations, busy_until=job_busy_until))
+
+        return jobs
 
     def _at_decision_step(self) -> bool:
         """Check if we're at a decision step (at least one job can be scheduled)."""
-        return len(self._get_ready_operations()) > 0
+        return len(self.state.available_jobs()) > 0
 
     def _validate_action(self, action: JSSPAction) -> bool:
         """Validate that an action is legal."""
@@ -204,15 +202,13 @@ class JSSPEnvironment(Environment):
             for i in range(self.nb_machines)
         ]
 
-        ready_ops = self._get_ready_operations()
+        jobs = self._get_jobs()
 
         return JSSPObservation(
             done=self.completed_jobs >= len(self.jobs),
             episode_id=self.episode_id,
             step_count=self.step_count,
             machines=machines,
-            ready_operations=ready_ops,
-            completed_jobs=self.completed_jobs,
-            total_jobs=len(self.jobs),
+            jobs=jobs,
             reward=0.0,  # Default, overwritten in step()
         )

src/jssp_openenv/solver.py

@@ -13,7 +13,7 @@ def solve_jssp(
 
     while not result.done:
         if verbose:
-            print(f"Step {obs.step_count}: {obs.ready_operations}")
+            print(f"Step {obs.step_count}: {', '.join([str(job.job_id) for job in obs.available_jobs()])}")
         action = policy.act(obs)
         if verbose:
             print(f"Action: {action}")
@@ -21,13 +21,13 @@ def solve_jssp(
         # Record scheduled events
         if action.job_ids:
             for job_id in action.job_ids:
-                operation = next((op for op in obs.ready_operations if op.job_id == job_id), None)
-                assert operation is not None
+                job = next((job for job in obs.available_jobs() if job.job_id == job_id), None)
+                assert job is not None
                 event = ScheduledEvent(
                     job_id=job_id,
-                    machine_id=operation.machine_id,
+                    machine_id=job.operations[0][0],
                     start_time=obs.step_count,
-                    end_time=obs.step_count + operation.duration,
+                    end_time=obs.step_count + job.operations[0][1],
                 )
                 scheduled_events.append(event)