# usail-hkust/llmtscs architecture

> Architecture diagram of the GitHub repository usail-hkust/llmtscs, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-20). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/usail-hkust/llmtscs
- Repository: https://github.com/usail-hkust/llmtscs

## Overview

1. **What this repository is**
   LLMLight is a traffic-signal-control framework centered on a shared simulation environment. It compares three kinds of controllers: heuristic baselines, RL/DNN baselines, and LLM-based agents. The repo also contains a separate LightGPT fine-tuning pipeline for adapting an LLM to traffic signal decisions.

2. **Core runtime boundary: simulation vs. decision policy**
   The main architectural split is between the **traffic simulation layer** and the **control policy layer**. The simulation layer is built around CityFlow and the road-network/traffic scenario JSON files under `data/`. The policy layer contains all agents that observe the current traffic state and emit a phase/action for each intersection.

3. **Environment and state construction**
   `utils/cityflow_env.py` is the central environment wrapper. Supporting utilities in `utils/` handle sampling, replay construction, update loops, evaluation, and transfer testing. The `data/<city>/<scenario>/` folders provide the road network definition plus anonymized real/synthetic traffic traces that drive the simulator.

4. **Agent hierarchy**
   `models/agent.py` is the base abstraction for controllers. Everything else in `models/` is a concrete policy implementation: fixed-time, random, max-pressure variants, MPLight/CoLight/PressLight/AttendLight, plus advanced and efficient variants. This makes the model layer a plug-in set of interchangeable controllers over the same environment API.

5. **Heuristic and RL baselines**
   The root `run_*.py` scripts are the execution entrypoints for benchmark runs. Scripts such as `run_fixedtime.py`, `run_maxpressure.py`, `run_mplight.py`, `run_colight.py`, and their `advanced_` / `efficient_` counterparts wire a specific agent class to the environment, execute training or evaluation, and produce logs/results for comparison.

6. **LLM-based inference path**
   `run_chatgpt.py`, `run_open_LLM.py`, and `run_open_LLM_with_vllm.py` provide the LLM control path. Here, the simulator state is transformed into a textual prompt, sent to either a hosted model (ChatGPT) or a local open-source model, and the generated response is mapped back to a traffic-signal action. `models/chatgpt.py` holds the external API setup for OpenAI-style calls.

7. **Prompting and instruction design**
   The `prompts/` directory stores prompt templates and domain-knowledge material used to encode traffic state into natural-language instructions. `models/chatgpt.py` and the utilities around it sit at the boundary between structured simulator state and the LLM’s text interface.

8. **LightGPT fine-tuning pipeline**
   The `finetune/` package is a separate training subsystem for adapting a base LLM into LightGPT. `run_imitation_finetune.py` performs supervised imitation learning from collected traffic decision data, `merge_lora.py` merges LoRA adapters back into a base model, and `run_policy_refinement.py` performs critic-guided refinement after data collection. `state_action_2_instructions.py` converts simulator state/action pairs into instruction-style training examples.

9. **Data collection for refinement**
   `run_policy_refinement_data_collection.py` bridges RL and LLM training: it uses a trained RL controller as a teacher to generate policy-refinement data, which is then stored under `data/cgpr/`. This creates a two-stage flow: simulator rollout from an RL agent, then instruction-formatted training data for the LLM adapter.

10. **Evaluation and analysis utilities**
   `utils/model_test.py`, `utils/RL_transfer_test.py`, and related helpers support benchmarking and cross-scenario transfer. `utils/llm_aft_trainer.py` and ranking/loss utilities suggest additional training logic for alignment or preference-style optimization, but they remain within the same simulator-to-policy pipeline.

11. **Frontend replay/visualization**
   `frontend/page/` is a static browser-based viewer used to replay and inspect traffic-control behavior. It is not the control runtime itself; it consumes replay artifacts so users can visualize intersection states and agent decisions after an experiment.

12. **Repository-level artifacts**
   The `data/` tree is the main scenario asset store, while `media/` holds figures and demo materials. The README also references model checkpoints and logs as external outputs of the runs. Architecturally, those artifacts are downstream products of the training/evaluation scripts rather than first-class application services.

13. **Important external dependencies**
   CityFlow is the simulation engine and requires Linux. TensorFlow appears to support older RL components, while Hugging Face `transformers`, `peft`, `accelerate`, `datasets`, and optionally `vllm` power the LLM side. `wandb` is used for experiment tracking. These dependencies matter because they define two distinct execution modes: classic RL/simulation and LLM inference/fine-tuning.

Main entry points: [`models/agent.py`](https://github.com/usail-hkust/llmtscs/blob/main/models/agent.py) (Agent base), [`utils/cityflow_env.py`](https://github.com/usail-hkust/llmtscs/blob/main/utils/cityflow_env.py) (CityFlow env), [`run_chatgpt.py`](https://github.com/usail-hkust/llmtscs/blob/main/run_chatgpt.py) (Run ChatGPT), [`finetune/run_policy_refinement_data_collection.py`](https://github.com/usail-hkust/llmtscs/blob/main/finetune/run_policy_refinement_data_collection.py) (Collect data)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_runtime["Simulation runtime"]
  node_utils_cityflow_env["CityFlow env<br/>env wrapper<br/>[cityflow_env.py]"]
  node_utils_construct_sample["Samples<br/>data sampler"]
  node_utils_updater["Updater<br/>rollout loop<br/>[updater.py]"]
  node_utils_pipeline["Pipeline<br/>execution flow<br/>[pipeline.py]"]
  node_utils_generator["Generator<br/>[generator.py]"]
end

subgraph group_policies["Control policies"]
  node_models_agent["Agent base<br/>policy interface<br/>[agent.py]"]
  node_models_fixedtime_agent["Fixed-time<br/>heuristic policy<br/>[fixedtime_agent.py]"]
  node_models_maxpressure_agent["Max-pressure<br/>heuristic policy"]
  node_models_mplight_agent["MPLight<br/>RL policy<br/>[mplight_agent.py]"]
  node_models_colight_agent["CoLight<br/>RL policy<br/>[colight_agent.py]"]
  node_models_presslight_one["PressLight<br/>RL policy<br/>[presslight_one.py]"]
  node_models_chatgpt["ChatGPT API<br/>llm adapter<br/>[chatgpt.py]"]
  node_run_fixedtime["Run fixed<br/>benchmark entrypoint<br/>[run_fixedtime.py]"]
  node_run_mplight["Run MPLight<br/>benchmark entrypoint<br/>[run_mplight.py]"]
  node_run_chatgpt["Run ChatGPT<br/>llm inference<br/>[run_chatgpt.py]"]
  node_run_open_llm["Run open LLM<br/>llm inference<br/>[run_open_LLM.py]"]
  node_prompts_domain["Domain prompts<br/>prompt templates"]
end

subgraph group_finetune["LightGPT training"]
  node_finetune_state_action["State-action<br/>instruction builder"]
  node_finetune_collect["Collect data<br/>data collection"]
  node_finetune_imitation["Imitation tune<br/>supervised finetune"]
  node_finetune_refine["Policy refine<br/>critic refinement"]
  node_finetune_merge["Merge LoRA<br/>adapter merge<br/>[merge_lora.py]"]
end

subgraph group_tools["Support tools"]
  node_utils_rl_transfer["Transfer test<br/>evaluation"]
  node_utils_model_test["Model test<br/>evaluation<br/>[model_test.py]"]
  node_utils_aft_rank_loss["Rank loss<br/>preference loss"]
end

subgraph group_ui["Replay viewer"]
  node_frontend_page["Replay viewer<br/>browser ui<br/>[index.html]"]
end

node_run_fixedtime -->|"uses"| node_models_fixedtime_agent
node_run_fixedtime -->|"simulates"| node_utils_cityflow_env
node_run_mplight -->|"uses"| node_models_mplight_agent
node_run_mplight -->|"simulates"| node_utils_cityflow_env
node_run_chatgpt -->|"calls"| node_models_chatgpt
node_run_chatgpt -->|"prompts with"| node_prompts_domain
node_run_chatgpt -->|"observes"| node_utils_cityflow_env
node_run_open_llm -->|"prompts with"| node_prompts_domain
node_run_open_llm -->|"observes"| node_utils_cityflow_env
node_utils_cityflow_env -->|"feeds"| node_utils_construct_sample
node_utils_construct_sample -->|"drives"| node_utils_updater
node_utils_updater -->|"orchestrated by"| node_utils_pipeline
node_utils_pipeline -->|"prepares"| node_utils_generator
node_finetune_collect -->|"teacher"| node_models_mplight_agent
node_finetune_collect -->|"formats"| node_finetune_state_action
node_finetune_state_action -->|"trains"| node_finetune_imitation
node_finetune_imitation -->|"exports"| node_finetune_merge
node_finetune_collect -->|"supplies"| node_finetune_refine
node_finetune_refine -->|"uses"| node_utils_aft_rank_loss
node_utils_model_test -->|"supports"| node_utils_rl_transfer
node_frontend_page -->|"views"| node_utils_model_test
node_models_agent -->|"base for"| node_models_fixedtime_agent
node_models_agent -->|"base for"| node_models_maxpressure_agent
node_models_agent -->|"base for"| node_models_mplight_agent
node_models_agent -->|"base for"| node_models_colight_agent
node_models_agent -->|"base for"| node_models_presslight_one

click node_utils_cityflow_env "https://github.com/usail-hkust/llmtscs/blob/main/utils/cityflow_env.py"
click node_utils_construct_sample "https://github.com/usail-hkust/llmtscs/blob/main/utils/construct_sample.py"
click node_utils_updater "https://github.com/usail-hkust/llmtscs/blob/main/utils/updater.py"
click node_utils_pipeline "https://github.com/usail-hkust/llmtscs/blob/main/utils/pipeline.py"
click node_utils_generator "https://github.com/usail-hkust/llmtscs/blob/main/utils/generator.py"
click node_models_agent "https://github.com/usail-hkust/llmtscs/blob/main/models/agent.py"
click node_models_fixedtime_agent "https://github.com/usail-hkust/llmtscs/blob/main/models/fixedtime_agent.py"
click node_models_maxpressure_agent "https://github.com/usail-hkust/llmtscs/blob/main/models/maxpressure_agent.py"
click node_models_mplight_agent "https://github.com/usail-hkust/llmtscs/blob/main/models/mplight_agent.py"
click node_models_colight_agent "https://github.com/usail-hkust/llmtscs/blob/main/models/colight_agent.py"
click node_models_presslight_one "https://github.com/usail-hkust/llmtscs/blob/main/models/presslight_one.py"
click node_models_chatgpt "https://github.com/usail-hkust/llmtscs/blob/main/models/chatgpt.py"
click node_run_fixedtime "https://github.com/usail-hkust/llmtscs/blob/main/run_fixedtime.py"
click node_run_mplight "https://github.com/usail-hkust/llmtscs/blob/main/run_mplight.py"
click node_run_chatgpt "https://github.com/usail-hkust/llmtscs/blob/main/run_chatgpt.py"
click node_run_open_llm "https://github.com/usail-hkust/llmtscs/blob/main/run_open_LLM.py"
click node_prompts_domain "https://github.com/usail-hkust/llmtscs/blob/main/prompts/prompt_domain_knowledge.json"
click node_finetune_state_action "https://github.com/usail-hkust/llmtscs/blob/main/finetune/state_action_2_instructions.py"
click node_finetune_collect "https://github.com/usail-hkust/llmtscs/blob/main/finetune/run_policy_refinement_data_collection.py"
click node_finetune_imitation "https://github.com/usail-hkust/llmtscs/blob/main/finetune/run_imitation_finetune.py"
click node_finetune_refine "https://github.com/usail-hkust/llmtscs/blob/main/finetune/run_policy_refinement.py"
click node_finetune_merge "https://github.com/usail-hkust/llmtscs/blob/main/finetune/merge_lora.py"
click node_utils_rl_transfer "https://github.com/usail-hkust/llmtscs/blob/main/utils/RL_transfer_test.py"
click node_utils_model_test "https://github.com/usail-hkust/llmtscs/blob/main/utils/model_test.py"
click node_utils_aft_rank_loss "https://github.com/usail-hkust/llmtscs/blob/main/utils/aft_rank_loss_utils.py"
click node_frontend_page "https://github.com/usail-hkust/llmtscs/blob/main/frontend/page/index.html"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_utils_cityflow_env,node_utils_construct_sample,node_utils_updater,node_utils_pipeline,node_utils_generator toneBlue
class node_models_agent,node_models_fixedtime_agent,node_models_maxpressure_agent,node_models_mplight_agent,node_models_colight_agent,node_models_presslight_one,node_models_chatgpt,node_run_fixedtime,node_run_mplight,node_run_chatgpt,node_run_open_llm,node_prompts_domain toneAmber
class node_finetune_state_action,node_finetune_collect,node_finetune_imitation,node_finetune_refine,node_finetune_merge toneMint
class node_utils_rl_transfer,node_utils_model_test,node_utils_aft_rank_loss toneRose
class node_frontend_page toneIndigo
```

## Components

### Simulation runtime

CityFlow environment, data assets, and shared rollout/evaluation utilities.

- [CityFlow env](https://github.com/usail-hkust/llmtscs/blob/main/utils/cityflow_env.py): env wrapper · Central simulator wrapper that exposes state, action, and step APIs. · `utils/cityflow_env.py`
- [Samples](https://github.com/usail-hkust/llmtscs/blob/main/utils/construct_sample.py): data sampler · Builds training and rollout samples from simulator traces. · `utils/construct_sample.py`
- [Updater](https://github.com/usail-hkust/llmtscs/blob/main/utils/updater.py): rollout loop · Coordinates repeated environment updates during an experiment. · `utils/updater.py`
- [Pipeline](https://github.com/usail-hkust/llmtscs/blob/main/utils/pipeline.py): execution flow · Orchestrates data flow between environment, sampling, and evaluation stages. · `utils/pipeline.py`
- [Generator](https://github.com/usail-hkust/llmtscs/blob/main/utils/generator.py): scenario generator · Creates or prepares traffic scenarios and related artifacts. · `utils/generator.py`

### Control policies

Heuristic, RL, and LLM-based controllers that plug into the same environment API.

- [Agent base](https://github.com/usail-hkust/llmtscs/blob/main/models/agent.py): policy interface · Common controller abstraction implemented by all policy types. · `models/agent.py`
- [Fixed-time](https://github.com/usail-hkust/llmtscs/blob/main/models/fixedtime_agent.py): heuristic policy · Static signal timing baseline. · `models/fixedtime_agent.py`
- [Max-pressure](https://github.com/usail-hkust/llmtscs/blob/main/models/maxpressure_agent.py): heuristic policy · Pressure-based traffic-signal controller baseline. · `models/maxpressure_agent.py`
- [MPLight](https://github.com/usail-hkust/llmtscs/blob/main/models/mplight_agent.py): RL policy · Graph-based reinforcement-learning controller. · `models/mplight_agent.py`
- [CoLight](https://github.com/usail-hkust/llmtscs/blob/main/models/colight_agent.py): RL policy · Cooperative multi-intersection learning policy. · `models/colight_agent.py`
- [PressLight](https://github.com/usail-hkust/llmtscs/blob/main/models/presslight_one.py): RL policy · Pressure-aware deep controller variant. · `models/presslight_one.py`
- [ChatGPT API](https://github.com/usail-hkust/llmtscs/blob/main/models/chatgpt.py): llm adapter · External LLM call boundary and response handling. · `models/chatgpt.py`
- [Run fixed](https://github.com/usail-hkust/llmtscs/blob/main/run_fixedtime.py): benchmark entrypoint · Executes the fixed-time baseline in the shared simulator. · `run_fixedtime.py`
- [Run MPLight](https://github.com/usail-hkust/llmtscs/blob/main/run_mplight.py): benchmark entrypoint · Runs MPLight training or evaluation against CityFlow. · `run_mplight.py`
- [Run ChatGPT](https://github.com/usail-hkust/llmtscs/blob/main/run_chatgpt.py): llm inference · Turns simulator state into prompts and maps LLM output back to actions. · `run_chatgpt.py`
- [Run open LLM](https://github.com/usail-hkust/llmtscs/blob/main/run_open_LLM.py): llm inference · Local open-model control path for traffic-signal decisions. · `run_open_LLM.py`
- [Domain prompts](https://github.com/usail-hkust/llmtscs/blob/main/prompts/prompt_domain_knowledge.json): prompt templates · Traffic-domain knowledge and instruction templates for LLM prompting. · `prompts/prompt_domain_knowledge.json`

### LightGPT training

Separate pipeline for instruction data creation, imitation tuning, and refinement.

- [State-action](https://github.com/usail-hkust/llmtscs/blob/main/finetune/state_action_2_instructions.py): instruction builder · Converts simulator state/action pairs into instruction-format examples. · `finetune/state_action_2_instructions.py`
- [Collect data](https://github.com/usail-hkust/llmtscs/blob/main/finetune/run_policy_refinement_data_collection.py): data collection · Uses an RL teacher to gather policy-refinement training data. · `finetune/run_policy_refinement_data_collection.py`
- [Imitation tune](https://github.com/usail-hkust/llmtscs/blob/main/finetune/run_imitation_finetune.py): supervised finetune · Supervised LightGPT adaptation from collected traffic decisions. · `finetune/run_imitation_finetune.py`
- [Policy refine](https://github.com/usail-hkust/llmtscs/blob/main/finetune/run_policy_refinement.py): critic refinement · Post-collection refinement stage for LightGPT-style policies. · `finetune/run_policy_refinement.py`
- [Merge LoRA](https://github.com/usail-hkust/llmtscs/blob/main/finetune/merge_lora.py): adapter merge · Merges LoRA adapters into a base model after tuning. · `finetune/merge_lora.py`

### Support tools

Transfer testing, ranking losses, and replay/visualization utilities.

- [Transfer test](https://github.com/usail-hkust/llmtscs/blob/main/utils/RL_transfer_test.py): evaluation · Cross-scenario transfer benchmarking for trained controllers. · `utils/RL_transfer_test.py`
- [Model test](https://github.com/usail-hkust/llmtscs/blob/main/utils/model_test.py): evaluation · Shared benchmarking and result-checking utility. · `utils/model_test.py`
- [Rank loss](https://github.com/usail-hkust/llmtscs/blob/main/utils/aft_rank_loss_utils.py): preference loss · Ranking and alignment loss helpers used in advanced training. · `utils/aft_rank_loss_utils.py`

### Replay viewer

Browser-based inspection of traffic experiments and replays.

- [Replay viewer](https://github.com/usail-hkust/llmtscs/blob/main/frontend/page/index.html): browser ui · Static frontend for inspecting traffic experiment replays. · `frontend/page/index.html`

## Connections

- Run fixed → Fixed-time: uses · Entry script wires a specific controller to the environment.
- Run fixed → CityFlow env: simulates · Runs the baseline inside the shared CityFlow wrapper.
- Run MPLight → MPLight: uses · Benchmark entrypoint for the RL controller.
- Run MPLight → CityFlow env: simulates · Shares the same environment API as other controllers.
- Run ChatGPT → ChatGPT API: calls · Uses the OpenAI-style API boundary for inference.
- Run ChatGPT → Domain prompts: prompts with · Loads domain knowledge and instruction templates.
- Run ChatGPT → CityFlow env: observes · Converts simulator state into text prompts and actions.
- Run open LLM → Domain prompts: prompts with · Uses the same prompt material for local model inference.
- Run open LLM → CityFlow env: observes · Maps simulator observations into LLM actions.
- CityFlow env → Samples: feeds · Produces samples from simulation traces and rollouts.
- Samples → Updater: drives · Sampling supports repeated rollout/update loops.
- Updater → Pipeline: orchestrated by · Higher-level pipeline coordinates execution stages.
- Pipeline → Generator: prepares · Pipeline can create or transform scenario artifacts.
- Collect data → MPLight: teacher · Uses a trained RL controller to generate refinement data.
- Collect data → State-action: formats · Converts rollouts into instruction-style training pairs.
- State-action → Imitation tune: trains · Feeds supervised LightGPT imitation learning.
- Imitation tune → Merge LoRA: exports · Adapters are merged back into a base model after tuning.
- Collect data → Policy refine: supplies · Collected data is reused for refinement.
- Policy refine → Rank loss: uses · Alignment/refinement can rely on ranking losses.
- Model test → Transfer test: supports · Shared evaluation utilities underpin transfer testing.
- Replay viewer → Model test: views · Replay viewer consumes experiment outputs from runs.
- Agent base → Fixed-time: base for · Common policy interface implemented by controllers.
- Agent base → Max-pressure: base for · Common policy interface implemented by controllers.
- Agent base → MPLight: base for · Common policy interface implemented by controllers.
- Agent base → CoLight: base for · Common policy interface implemented by controllers.
- Agent base → PressLight: base for · Common policy interface implemented by controllers.