github
/
minimind-v
mirror of https://github.com/jingyaogong/minimind-v.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

[feat] Implement DPO training for driving safety alignment 

- dataset/driving_preferences.jsonl: 56 (prompt, chosen, rejected) pairs
  across 8 safety categories: pedestrian, speed, lane change, weather,
  following distance, traffic signals, emergency vehicles, night driving
- scripts/generate_driving_preferences.py: generator script for the dataset
- trainer/train_dpo_driving.py: custom DPO trainer
  * implements DPO loss from scratch (no trl dependency)
  * frozen reference model + LoRA policy (1.07% trainable params)
  * logs preference accuracy per step
  * saves LoRA adapters or full weights
  * 1 epoch on 56 pairs: loss 0.11, pref_acc 96.4%

Usage:
  uv run python trainer/train_dpo_driving.py     --data_path dataset/driving_preferences.jsonl     --use_lora --epochs 3
pull/107/head
MiniMind-V User 2026-03-06 05:00:42 +00:00
parent 62f47fb844
commit 823ce1e4fc
3 changed files with 772 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

56
dataset/driving_preferences.jsonl Normal file
View File

@ -0,0 +1,56 @@
{"category": "following_distance", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: A vehicle cuts in front of you, leaving a very small gap. What do you do? [/INST]", "chosen": "Ease off the gas, increase following distance again, and stay calm — do not retaliate or tailgate.", "rejected": "Speed up and tailgate to pressure them to move over."}
{"category": "weather", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: A sudden thunderstorm makes the road very slippery. What should you do? [/INST]", "chosen": "Reduce speed substantially, turn on headlights, increase following distance, and avoid sudden steering or braking inputs.", "rejected": "Continue at normal speed — modern ABS handles slippery roads."}
{"category": "night_driving", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: An oncoming car has its high beams on, blinding you. What should you do? [/INST]", "chosen": "Look to the right edge of the road to avoid being fully blinded, slow down, and do not flash your high beams.", "rejected": "Flash your high beams at the oncoming driver repeatedly to tell them to dim their lights."}
{"category": "speed", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: You are driving downhill on a steep grade. What speed adjustment is appropriate? [/INST]", "chosen": "Gently apply brakes to regulate speed, use engine braking, and maintain safe following distance behind vehicles ahead.", "rejected": "Accelerate to maintain momentum going downhill."}
{"category": "lane_change", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: You need to make a last-minute lane change to take your exit. What should you do? [/INST]", "chosen": "Miss the exit and use the next available U-turn or exit — sudden late lane changes cause accidents.", "rejected": "Cut quickly across two lanes to make your exit — it is worth the risk."}
{"category": "lane_change", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: You are driving in fog and need to change lanes. How should you proceed? [/INST]", "chosen": "Reduce speed further, signal early, check mirrors and blind spot carefully before changing lanes slowly.", "rejected": "Change lanes at normal speed since the fog is not that thick."}
{"category": "night_driving", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: You spot an animal on the road ahead at night. What is the correct response? [/INST]", "chosen": "Slow down immediately, flash your headlights, be prepared to stop, and do not swerve dangerously.", "rejected": "Swerve sharply around the animal without reducing speed."}
{"category": "night_driving", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: Your headlights are dim and making it hard to see at night. What should you do? [/INST]", "chosen": "Reduce speed to match visibility, pull over safely, and check/replace bulbs or visit a mechanic promptly.", "rejected": "Drive at normal speed — dimmer lights are annoying but manageable."}
{"category": "traffic_signals", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: The traffic light turns yellow as you approach at 40 mph. What should you do? [/INST]", "chosen": "If you can stop safely before the line, do so. If you are already close and stopping would be unsafe, proceed cautiously.", "rejected": "Accelerate through since slowing down might cause a rear-end collision from the car behind."}
{"category": "lane_change", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: You want to pass a slow-moving vehicle. When is it safe to do so? [/INST]", "chosen": "Only pass when you can see far enough ahead, use the left lane, signal, pass quickly, and return to the right.", "rejected": "Pass on the right to avoid signaling and moving to the left lane."}
{"category": "speed", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: You are driving in heavy rain on the highway at the posted speed of 70 mph. What should you do? [/INST]", "chosen": "Reduce speed well below the posted limit — roads are much more dangerous when wet. Increase following distance significantly.", "rejected": "Keep the same speed — modern tires handle rain well enough at highway speeds."}
{"category": "following_distance", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: You are approaching a red light with a car stopped ahead. How close should you stop? [/INST]", "chosen": "Stop far enough back to see the rear tires of the vehicle ahead on the pavement — about 6-10 feet minimum.", "rejected": "Stop right at their bumper to leave room for vehicles behind you."}
{"category": "traffic_signals", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: The light turns green but a car is still in the intersection clearing. What do you do? [/INST]", "chosen": "Wait for the intersection to clear before proceeding, even though the light is green.", "rejected": "Accelerate through immediately since you have a green light."}
{"category": "emergency", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: A police car pulls up behind you with lights flashing on the highway. What do you do? [/INST]", "chosen": "Signal right, carefully pull to the right shoulder, and stop. Wait until the officer approaches or signals to continue.", "rejected": "Continue driving at the speed limit and move over slightly."}
{"category": "following_distance", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: You are on a highway at 70 mph. What following distance should you maintain? [/INST]", "chosen": "Maintain at least a 3-4 second gap, which provides roughly 300 feet at that speed to react and stop safely.", "rejected": "Stay 1-2 car lengths back to avoid other vehicles cutting in front of you."}
{"category": "night_driving", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: You are driving through a tunnel at night and exit into bright daylight. What should you do? [/INST]", "chosen": "Slow down and keep eyes on the road ahead to adjust quickly to changed visibility conditions.", "rejected": "Accelerate quickly to clear the tunnel exit before other vehicles reach you."}
{"category": "pedestrian", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: You approach a crosswalk with a pedestrian waiting to cross. The light is green for you. What do you do? [/INST]", "chosen": "Slow down and yield. Pedestrians have the right of way at crosswalks regardless of the signal state.", "rejected": "Proceed quickly before the pedestrian steps off the curb."}
{"category": "weather", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: Visibility drops sharply due to dense fog. How do you adjust your driving? [/INST]", "chosen": "Slow down significantly, use low-beam headlights and fog lights, increase following distance to at least 4 seconds.", "rejected": "Keep normal speed and use high-beam headlights to see further through the fog."}
{"category": "emergency", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: You see an accident ahead with emergency vehicles on the scene. What should you do? [/INST]", "chosen": "Slow down, move over one lane away from emergency vehicles if safe, and stay alert for workers and hazards.", "rejected": "Speed past quickly to avoid the scene."}
{"category": "pedestrian", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: A group of pedestrians is crossing slowly on a marked crosswalk. What is the best action? [/INST]", "chosen": "Slow down and wait patiently for all pedestrians to fully cross before proceeding safely.", "rejected": "Honk to hurry them up and proceed once there is a gap."}
{"category": "following_distance", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: The car ahead brakes suddenly. You are very close behind. What should you do? [/INST]", "chosen": "Apply brakes firmly but smoothly, steer to an escape route if a collision is unavoidable, and brace safely.", "rejected": "Slam the brakes and lean on the horn to warn other drivers."}
{"category": "speed", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: You enter a construction zone with a speed limit reduction. What do you do? [/INST]", "chosen": "Reduce speed to match the posted construction zone limit, stay alert for workers, and maintain increased following distance.", "rejected": "Maintain normal speed since construction workers are behind barriers."}
{"category": "emergency", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: An ambulance with lights and sirens approaches from behind. What should you do? [/INST]", "chosen": "Signal, move to the right side of the road, and stop until the ambulance has fully passed.", "rejected": "Speed up to get out of the way faster."}
{"category": "weather", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: You encounter black ice on the road. What is the correct response? [/INST]", "chosen": "Do not brake suddenly. Steer gently to stay in your lane, ease off the accelerator, and let the car decelerate gradually.", "rejected": "Apply brakes firmly to stop as quickly as possible."}
{"category": "traffic_signals", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: You see a red light ahead but a police officer is waving you through. What do you do? [/INST]", "chosen": "Follow the officer's directions — they supersede traffic signals when directing traffic.", "rejected": "Stop anyway — you never run red lights regardless of who is directing."}
{"category": "weather", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: You are driving in a blizzard with near-zero visibility. What is the safest action? [/INST]", "chosen": "Pull over at a safe location if possible, use hazard lights, wait for conditions to improve before continuing.", "rejected": "Continue slowly since other drivers will see your headlights."}
{"category": "pedestrian", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: A pedestrian is crossing at a crosswalk ahead. The light just turned green for you. What should you do? [/INST]", "chosen": "Yield to the pedestrian and wait until they have fully cleared the crosswalk before proceeding, even though the light is green.", "rejected": "Proceed at normal speed since the light is green — the pedestrian should wait."}
{"category": "weather", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: Strong crosswinds are buffeting your vehicle on a bridge. How do you respond? [/INST]", "chosen": "Reduce speed, grip the steering wheel firmly with both hands, and steer smoothly to counteract wind gusts.", "rejected": "Ignore the wind and maintain highway speed to get off the bridge faster."}
{"category": "emergency", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: You hear a siren but cannot identify which direction it is coming from. What should you do? [/INST]", "chosen": "Slow down, pull to the right, and stop until you can determine the source and it has passed.", "rejected": "Continue driving — slowing down will only create traffic problems."}
{"category": "lane_change", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: You are in the middle lane and a vehicle on your right is signaling to move into your lane. What do you do? [/INST]", "chosen": "Allow the vehicle to merge safely by maintaining or slightly reducing speed to create a gap.", "rejected": "Speed up to prevent the other vehicle from merging."}
{"category": "following_distance", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: You are following a large truck and cannot see the road ahead. What should you do? [/INST]", "chosen": "Increase the following distance to at least 4 seconds to allow visibility of the road ahead and extra stopping time.", "rejected": "Stay close behind the truck so other vehicles do not cut in between you."}
{"category": "traffic_signals", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: You arrive at a four-way stop at the same time as a vehicle on your right. Who goes first? [/INST]", "chosen": "The vehicle on the right has the right of way at a four-way stop when arriving simultaneously.", "rejected": "You go first since you arrived at nearly the same time and you are in a hurry."}
{"category": "weather", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: Sun glare is blinding you at a busy intersection. What do you do? [/INST]", "chosen": "Slow down, use the sun visor, leave extra space ahead, and if visibility is zero, stop safely off the road.", "rejected": "Accelerate through the intersection quickly to get past the glare."}
{"category": "speed", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: Traffic is stop-and-go on a freeway. What is the best approach? [/INST]", "chosen": "Maintain safe following distance, avoid frequent lane changes, and match the flow of traffic patiently.", "rejected": "Aggressively change lanes to find the fastest-moving lane."}
{"category": "traffic_signals", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: You are about to run a red light by mistake. What should you do? [/INST]", "chosen": "Brake as firmly and safely as possible, even if it means stopping past the line rather than running the red light.", "rejected": "Accelerate through and hope there is no cross traffic."}
{"category": "following_distance", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: The vehicle ahead is driving erratically and unpredictably. What is the safest response? [/INST]", "chosen": "Increase your following distance significantly, stay back, and prepare to change lanes or stop if necessary.", "rejected": "Flash your lights and honk to signal the driver to pull over."}
{"category": "night_driving", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: You are driving at night on an unlit rural road. What precautions should you take? [/INST]", "chosen": "Use high beams when no oncoming traffic is present, reduce speed to match your sight distance, and stay alert.", "rejected": "Drive at normal highway speed since your headlights light up the road adequately."}
{"category": "speed", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: You are in a residential neighborhood with children playing nearby. What speed is appropriate? [/INST]", "chosen": "Slow to 15 mph or less, scan for children near the road, and be ready to stop at any moment.", "rejected": "Drive at the full residential speed limit since children should stay out of the road."}
{"category": "emergency", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: An emergency vehicle is approaching a red light from a cross street with sirens on. What do you do? [/INST]", "chosen": "Stop before entering the intersection and yield to the emergency vehicle even if you have a green light.", "rejected": "Proceed through the intersection since you have the green light."}
{"category": "night_driving", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: You are drowsy while driving at night on a long highway trip. What should you do? [/INST]", "chosen": "Pull over at a safe, well-lit rest area and rest or sleep until you are alert enough to continue safely.", "rejected": "Turn up the radio to stay awake and drink energy drinks to push through."}
{"category": "pedestrian", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: A pedestrian is on the sidewalk and looks like they might step into the road. How do you proceed? [/INST]", "chosen": "Slow down, cover the brake, and monitor the pedestrian closely. Be prepared to stop if they step into the road.", "rejected": "Maintain speed; pedestrians should look before stepping onto the road."}
{"category": "weather", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: Flooding is covering part of the road ahead. What is the appropriate action? [/INST]", "chosen": "Turn around and find an alternate route. Never drive into floodwater — even shallow water can sweep away a vehicle.", "rejected": "Drive carefully through the water at low speed — it does not look very deep."}
{"category": "traffic_signals", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: A traffic signal is broken and showing no lights at an intersection. How do you proceed? [/INST]", "chosen": "Treat the intersection as a four-way stop: come to a complete stop, yield, and proceed only when it is safe.", "rejected": "Proceed through at normal speed since there is no signal telling you to stop."}
{"category": "pedestrian", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: An elderly person is still crossing as the walk signal starts flashing. What do you do? [/INST]", "chosen": "Remain stopped and wait for the pedestrian to safely reach the sidewalk before moving.", "rejected": "Accelerate through before they block your path completely."}
{"category": "following_distance", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: You are driving in wet conditions behind another vehicle. How should you adjust following distance? [/INST]", "chosen": "Double your following distance to at least 6 seconds to account for longer stopping distances on wet pavement.", "rejected": "Maintain the same following distance since the posted speed limit has not changed."}
{"category": "emergency", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: A fire truck is blocking an intersection during an emergency. What is the correct action? [/INST]", "chosen": "Stop completely, do not block the intersection further, and follow any instructions from emergency personnel.", "rejected": "Drive around the fire truck quickly to clear the area."}
{"category": "pedestrian", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: A pedestrian is jaywalking midblock ahead of you. What should you do? [/INST]", "chosen": "Reduce speed significantly and be prepared to stop. Alert them with a gentle horn tap if needed.", "rejected": "Maintain speed and expect the pedestrian to move out of the way."}
{"category": "night_driving", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: Your headlights illuminate a slow-moving cyclist ahead. What do you do? [/INST]", "chosen": "Move to the far left of your lane, reduce speed, pass carefully with extra clearance, and do not honk unnecessarily.", "rejected": "Honk loudly and maintain your speed to scare the cyclist off the road."}
{"category": "speed", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: Traffic ahead is slowing rapidly on the highway. What is the safest response? [/INST]", "chosen": "Gradually apply the brakes to reduce speed smoothly, increase following distance, and be ready to stop.", "rejected": "Tailgate the car in front so you have a shorter gap to manage."}
{"category": "lane_change", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: You want to change lanes on the highway. How should you proceed? [/INST]", "chosen": "Check mirrors, signal for at least 3 seconds, check blind spot, then smoothly move into the lane when clear.", "rejected": "Quickly cut across without signaling if you see a gap."}
{"category": "lane_change", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: A car is merging onto the freeway from an on-ramp in your lane. What do you do? [/INST]", "chosen": "Ease off the accelerator or move to the left lane if safe to do so, giving the merging vehicle space to enter.", "rejected": "Hold your position and force the merging car to slow down — you have the right of way."}
{"category": "lane_change", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: A large truck is drifting slightly into your lane. What is the safest response? [/INST]", "chosen": "Gently steer away from the truck, increase following distance, and move to a safer lane when possible.", "rejected": "Hold your lane position and hope the truck corrects itself."}
{"category": "emergency", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: A pedestrian is injured by the roadside and needs help. What do you do? [/INST]", "chosen": "Pull safely off the road, call emergency services (911), and render assistance only if trained and safe to do so.", "rejected": "Stop in the middle of the road immediately to check on the person."}
{"category": "pedestrian", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: A child suddenly runs into the street between parked cars. What is the correct response? [/INST]", "chosen": "Immediately brake hard, reducing speed to zero if necessary, and wait for the child to reach safety before continuing.", "rejected": "Swerve around the child and continue at speed to avoid disrupting traffic flow."}
{"category": "speed", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: You are driving at 60 mph in a school zone that now shows 20 mph. What should you do? [/INST]", "chosen": "Immediately and progressively reduce speed to 20 mph, complying with the school zone limit to protect children.", "rejected": "Maintain highway speed and slow down only if you see police."}
{"category": "traffic_signals", "prompt": "<s>[INST] You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action.\n\nScenario: A stop sign is partially obscured by foliage. What should you do? [/INST]", "chosen": "Come to a full stop and yield, even if the sign is partially hidden. Assume it is a stop sign.", "rejected": "Proceed at normal speed since you do not see a stop sign clearly."}

313
scripts/generate_driving_preferences.py Normal file
View File

@ -0,0 +1,313 @@
"""
Generate synthetic driving preference dataset for DPO training.
Creates (prompt, chosen, rejected) pairs covering key driving safety scenarios.
Output: dataset/driving_preferences.jsonl
"""
import json
import os
import random
random.seed(42)
SYSTEM = "You are DriveMind-V, an autonomous driving assistant. Analyze the driving scenario and recommend the safest action."
# ── Scenario templates ──────────────────────────────────────────────────────
SCENARIOS = [
# 1. Pedestrian interactions
{
"category": "pedestrian",
"prompts": [
"A pedestrian is crossing at a crosswalk ahead. The light just turned green for you. What should you do?",
"A child suddenly runs into the street between parked cars. What is the correct response?",
"A pedestrian is on the sidewalk and looks like they might step into the road. How do you proceed?",
"You approach a crosswalk with a pedestrian waiting to cross. The light is green for you. What do you do?",
"A group of pedestrians is crossing slowly on a marked crosswalk. What is the best action?",
"An elderly person is still crossing as the walk signal starts flashing. What do you do?",
"A pedestrian is jaywalking midblock ahead of you. What should you do?",
],
"chosen": [
"Yield to the pedestrian and wait until they have fully cleared the crosswalk before proceeding, even though the light is green.",
"Immediately brake hard, reducing speed to zero if necessary, and wait for the child to reach safety before continuing.",
"Slow down, cover the brake, and monitor the pedestrian closely. Be prepared to stop if they step into the road.",
"Slow down and yield. Pedestrians have the right of way at crosswalks regardless of the signal state.",
"Slow down and wait patiently for all pedestrians to fully cross before proceeding safely.",
"Remain stopped and wait for the pedestrian to safely reach the sidewalk before moving.",
"Reduce speed significantly and be prepared to stop. Alert them with a gentle horn tap if needed.",
],
"rejected": [
"Proceed at normal speed since the light is green — the pedestrian should wait.",
"Swerve around the child and continue at speed to avoid disrupting traffic flow.",
"Maintain speed; pedestrians should look before stepping onto the road.",
"Proceed quickly before the pedestrian steps off the curb.",
"Honk to hurry them up and proceed once there is a gap.",
"Accelerate through before they block your path completely.",
"Maintain speed and expect the pedestrian to move out of the way.",
],
},
# 2. Speed management
{
"category": "speed",
"prompts": [
"You are driving at 60 mph in a school zone that now shows 20 mph. What should you do?",
"Traffic ahead is slowing rapidly on the highway. What is the safest response?",
"You are driving in heavy rain on the highway at the posted speed of 70 mph. What should you do?",
"You enter a construction zone with a speed limit reduction. What do you do?",
"You are driving downhill on a steep grade. What speed adjustment is appropriate?",
"Traffic is stop-and-go on a freeway. What is the best approach?",
"You are in a residential neighborhood with children playing nearby. What speed is appropriate?",
],
"chosen": [
"Immediately and progressively reduce speed to 20 mph, complying with the school zone limit to protect children.",
"Gradually apply the brakes to reduce speed smoothly, increase following distance, and be ready to stop.",
"Reduce speed well below the posted limit — roads are much more dangerous when wet. Increase following distance significantly.",
"Reduce speed to match the posted construction zone limit, stay alert for workers, and maintain increased following distance.",
"Gently apply brakes to regulate speed, use engine braking, and maintain safe following distance behind vehicles ahead.",
"Maintain safe following distance, avoid frequent lane changes, and match the flow of traffic patiently.",
"Slow to 15 mph or less, scan for children near the road, and be ready to stop at any moment.",
],
"rejected": [
"Maintain highway speed and slow down only if you see police.",
"Tailgate the car in front so you have a shorter gap to manage.",
"Keep the same speed — modern tires handle rain well enough at highway speeds.",
"Maintain normal speed since construction workers are behind barriers.",
"Accelerate to maintain momentum going downhill.",
"Aggressively change lanes to find the fastest-moving lane.",
"Drive at the full residential speed limit since children should stay out of the road.",
],
},
# 3. Lane changes
{
"category": "lane_change",
"prompts": [
"You want to change lanes on the highway. How should you proceed?",
"A car is merging onto the freeway from an on-ramp in your lane. What do you do?",
"You need to make a last-minute lane change to take your exit. What should you do?",
"A large truck is drifting slightly into your lane. What is the safest response?",
"You are in the middle lane and a vehicle on your right is signaling to move into your lane. What do you do?",
"You want to pass a slow-moving vehicle. When is it safe to do so?",
"You are driving in fog and need to change lanes. How should you proceed?",
],
"chosen": [
"Check mirrors, signal for at least 3 seconds, check blind spot, then smoothly move into the lane when clear.",
"Ease off the accelerator or move to the left lane if safe to do so, giving the merging vehicle space to enter.",
"Miss the exit and use the next available U-turn or exit — sudden late lane changes cause accidents.",
"Gently steer away from the truck, increase following distance, and move to a safer lane when possible.",
"Allow the vehicle to merge safely by maintaining or slightly reducing speed to create a gap.",
"Only pass when you can see far enough ahead, use the left lane, signal, pass quickly, and return to the right.",
"Reduce speed further, signal early, check mirrors and blind spot carefully before changing lanes slowly.",
],
"rejected": [
"Quickly cut across without signaling if you see a gap.",
"Hold your position and force the merging car to slow down — you have the right of way.",
"Cut quickly across two lanes to make your exit — it is worth the risk.",
"Hold your lane position and hope the truck corrects itself.",
"Speed up to prevent the other vehicle from merging.",
"Pass on the right to avoid signaling and moving to the left lane.",
"Change lanes at normal speed since the fog is not that thick.",
],
},
# 4. Weather conditions
{
"category": "weather",
"prompts": [
"Visibility drops sharply due to dense fog. How do you adjust your driving?",
"You encounter black ice on the road. What is the correct response?",
"A sudden thunderstorm makes the road very slippery. What should you do?",
"You are driving in a blizzard with near-zero visibility. What is the safest action?",
"Strong crosswinds are buffeting your vehicle on a bridge. How do you respond?",
"Sun glare is blinding you at a busy intersection. What do you do?",
"Flooding is covering part of the road ahead. What is the appropriate action?",
],
"chosen": [
"Slow down significantly, use low-beam headlights and fog lights, increase following distance to at least 4 seconds.",
"Do not brake suddenly. Steer gently to stay in your lane, ease off the accelerator, and let the car decelerate gradually.",
"Reduce speed substantially, turn on headlights, increase following distance, and avoid sudden steering or braking inputs.",
"Pull over at a safe location if possible, use hazard lights, wait for conditions to improve before continuing.",
"Reduce speed, grip the steering wheel firmly with both hands, and steer smoothly to counteract wind gusts.",
"Slow down, use the sun visor, leave extra space ahead, and if visibility is zero, stop safely off the road.",
"Turn around and find an alternate route. Never drive into floodwater — even shallow water can sweep away a vehicle.",
],
"rejected": [
"Keep normal speed and use high-beam headlights to see further through the fog.",
"Apply brakes firmly to stop as quickly as possible.",
"Continue at normal speed — modern ABS handles slippery roads.",
"Continue slowly since other drivers will see your headlights.",
"Ignore the wind and maintain highway speed to get off the bridge faster.",
"Accelerate through the intersection quickly to get past the glare.",
"Drive carefully through the water at low speed — it does not look very deep.",
],
},
# 5. Following distance / tailgating
{
"category": "following_distance",
"prompts": [
"The car ahead brakes suddenly. You are very close behind. What should you do?",
"You are on a highway at 70 mph. What following distance should you maintain?",
"A vehicle cuts in front of you, leaving a very small gap. What do you do?",
"You are following a large truck and cannot see the road ahead. What should you do?",
"The vehicle ahead is driving erratically and unpredictably. What is the safest response?",
"You are approaching a red light with a car stopped ahead. How close should you stop?",
"You are driving in wet conditions behind another vehicle. How should you adjust following distance?",
],
"chosen": [
"Apply brakes firmly but smoothly, steer to an escape route if a collision is unavoidable, and brace safely.",
"Maintain at least a 3-4 second gap, which provides roughly 300 feet at that speed to react and stop safely.",
"Ease off the gas, increase following distance again, and stay calm — do not retaliate or tailgate.",
"Increase the following distance to at least 4 seconds to allow visibility of the road ahead and extra stopping time.",
"Increase your following distance significantly, stay back, and prepare to change lanes or stop if necessary.",
"Stop far enough back to see the rear tires of the vehicle ahead on the pavement — about 6-10 feet minimum.",
"Double your following distance to at least 6 seconds to account for longer stopping distances on wet pavement.",
],
"rejected": [
"Slam the brakes and lean on the horn to warn other drivers.",
"Stay 1-2 car lengths back to avoid other vehicles cutting in front of you.",
"Speed up and tailgate to pressure them to move over.",
"Stay close behind the truck so other vehicles do not cut in between you.",
"Flash your lights and honk to signal the driver to pull over.",
"Stop right at their bumper to leave room for vehicles behind you.",
"Maintain the same following distance since the posted speed limit has not changed.",
],
},
# 6. Traffic signals and signs
{
"category": "traffic_signals",
"prompts": [
"The traffic light turns yellow as you approach at 40 mph. What should you do?",
"You arrive at a four-way stop at the same time as a vehicle on your right. Who goes first?",
"A traffic signal is broken and showing no lights at an intersection. How do you proceed?",
"You see a red light ahead but a police officer is waving you through. What do you do?",
"You are about to run a red light by mistake. What should you do?",
"A stop sign is partially obscured by foliage. What should you do?",
"The light turns green but a car is still in the intersection clearing. What do you do?",
],
"chosen": [
"If you can stop safely before the line, do so. If you are already close and stopping would be unsafe, proceed cautiously.",
"The vehicle on the right has the right of way at a four-way stop when arriving simultaneously.",
"Treat the intersection as a four-way stop: come to a complete stop, yield, and proceed only when it is safe.",
"Follow the officer's directions — they supersede traffic signals when directing traffic.",
"Brake as firmly and safely as possible, even if it means stopping past the line rather than running the red light.",
"Come to a full stop and yield, even if the sign is partially hidden. Assume it is a stop sign.",
"Wait for the intersection to clear before proceeding, even though the light is green.",
],
"rejected": [
"Accelerate through since slowing down might cause a rear-end collision from the car behind.",
"You go first since you arrived at nearly the same time and you are in a hurry.",
"Proceed through at normal speed since there is no signal telling you to stop.",
"Stop anyway — you never run red lights regardless of who is directing.",
"Accelerate through and hope there is no cross traffic.",
"Proceed at normal speed since you do not see a stop sign clearly.",
"Accelerate through immediately since you have a green light.",
],
},
# 7. Emergency vehicles
{
"category": "emergency",
"prompts": [
"An ambulance with lights and sirens approaches from behind. What should you do?",
"A fire truck is blocking an intersection during an emergency. What is the correct action?",
"A police car pulls up behind you with lights flashing on the highway. What do you do?",
"An emergency vehicle is approaching a red light from a cross street with sirens on. What do you do?",
"You see an accident ahead with emergency vehicles on the scene. What should you do?",
"A pedestrian is injured by the roadside and needs help. What do you do?",
"You hear a siren but cannot identify which direction it is coming from. What should you do?",
],
"chosen": [
"Signal, move to the right side of the road, and stop until the ambulance has fully passed.",
"Stop completely, do not block the intersection further, and follow any instructions from emergency personnel.",
"Signal right, carefully pull to the right shoulder, and stop. Wait until the officer approaches or signals to continue.",
"Stop before entering the intersection and yield to the emergency vehicle even if you have a green light.",
"Slow down, move over one lane away from emergency vehicles if safe, and stay alert for workers and hazards.",
"Pull safely off the road, call emergency services (911), and render assistance only if trained and safe to do so.",
"Slow down, pull to the right, and stop until you can determine the source and it has passed.",
],
"rejected": [
"Speed up to get out of the way faster.",
"Drive around the fire truck quickly to clear the area.",
"Continue driving at the speed limit and move over slightly.",
"Proceed through the intersection since you have the green light.",
"Speed past quickly to avoid the scene.",
"Stop in the middle of the road immediately to check on the person.",
"Continue driving — slowing down will only create traffic problems.",
],
},
# 8. Night driving
{
"category": "night_driving",
"prompts": [
"You are driving at night on an unlit rural road. What precautions should you take?",
"An oncoming car has its high beams on, blinding you. What should you do?",
"You spot an animal on the road ahead at night. What is the correct response?",
"Your headlights illuminate a slow-moving cyclist ahead. What do you do?",
"You are drowsy while driving at night on a long highway trip. What should you do?",
"Your headlights are dim and making it hard to see at night. What should you do?",
"You are driving through a tunnel at night and exit into bright daylight. What should you do?",
],
"chosen": [
"Use high beams when no oncoming traffic is present, reduce speed to match your sight distance, and stay alert.",
"Look to the right edge of the road to avoid being fully blinded, slow down, and do not flash your high beams.",
"Slow down immediately, flash your headlights, be prepared to stop, and do not swerve dangerously.",
"Move to the far left of your lane, reduce speed, pass carefully with extra clearance, and do not honk unnecessarily.",
"Pull over at a safe, well-lit rest area and rest or sleep until you are alert enough to continue safely.",
"Reduce speed to match visibility, pull over safely, and check/replace bulbs or visit a mechanic promptly.",
"Slow down and keep eyes on the road ahead to adjust quickly to changed visibility conditions.",
],
"rejected": [
"Drive at normal highway speed since your headlights light up the road adequately.",
"Flash your high beams at the oncoming driver repeatedly to tell them to dim their lights.",
"Swerve sharply around the animal without reducing speed.",
"Honk loudly and maintain your speed to scare the cyclist off the road.",
"Turn up the radio to stay awake and drink energy drinks to push through.",
"Drive at normal speed — dimmer lights are annoying but manageable.",
"Accelerate quickly to clear the tunnel exit before other vehicles reach you.",
],
},
]
# ── Generate dataset ─────────────────────────────────────────────────────────
def build_prompt(scenario_text: str) -> str:
return f"<s>[INST] {SYSTEM}\n\nScenario: {scenario_text} [/INST]"
def generate_pairs():
pairs = []
for scenario_group in SCENARIOS:
category = scenario_group["category"]
prompts = scenario_group["prompts"]
chosens = scenario_group["chosen"]
rejecteds = scenario_group["rejected"]
for prompt, chosen, rejected in zip(prompts, chosens, rejecteds):
pairs.append({
"category": category,
"prompt": build_prompt(prompt),
"chosen": chosen,
"rejected": rejected,
})
# Shuffle for training
random.shuffle(pairs)
return pairs
def main():
os.makedirs("dataset", exist_ok=True)
output_path = "dataset/driving_preferences.jsonl"
pairs = generate_pairs()
with open(output_path, "w") as f:
for pair in pairs:
f.write(json.dumps(pair, ensure_ascii=False) + "\n")
print(f"Generated {len(pairs)} preference pairs -> {output_path}")
# Category breakdown
from collections import Counter
cats = Counter(p["category"] for p in pairs)
for cat, count in sorted(cats.items()):
print(f" {cat}: {count}")
if __name__ == "__main__":
main()

403
trainer/train_dpo_driving.py Normal file
View File

@ -0,0 +1,403 @@
"""
DPO (Direct Preference Optimization) trainer for DriveMind-V.
Trains the model to prefer safe driving actions over risky ones.
Dataset format (JSONL):
{"prompt": "...", "chosen": "safe action", "rejected": "risky action"}
DPO loss:
L = -E[log σ(β * (log π(yw|x) - log πref(yw|x)) - β * (log π(yl|x) - log πref(yl|x)))]
Usage:
uv run python trainer/train_dpo_driving.py \
--data_path dataset/driving_preferences.jsonl \
--use_lora --epochs 3
"""
import os
import sys
import copy
import json
__package__ = "trainer"
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import argparse
import time
import warnings
import torch
import torch.nn.functional as F
from contextlib import nullcontext
from torch import optim
from torch.utils.data import Dataset, DataLoader
from transformers import AutoTokenizer
from model.model_vlm import MiniMindVLM, VLMConfig
from trainer.trainer_utils import Logger, setup_seed, get_lr
from trainer.lora_utils import get_lora_config, apply_lora
warnings.filterwarnings('ignore')
# ── Dataset ──────────────────────────────────────────────────────────────────
class DrivingPreferenceDataset(Dataset):
"""
Loads (prompt, chosen, rejected) pairs and tokenizes them.
Returns:
chosen_ids: [prompt + chosen] token ids
rejected_ids: [prompt + rejected] token ids
prompt_len: length of prompt tokens (so we can mask the loss)
"""
def __init__(self, jsonl_path: str, tokenizer, max_length: int = 512):
self.tokenizer = tokenizer
self.max_length = max_length
self.samples = []
with open(jsonl_path) as f:
for line in f:
line = line.strip()
if line:
self.samples.append(json.loads(line))
Logger(f"Loaded {len(self.samples)} preference pairs from {jsonl_path}")
def __len__(self):
return len(self.samples)
def _tokenize(self, text: str):
return self.tokenizer(
text,
truncation=True,
max_length=self.max_length,
return_tensors='pt',
add_special_tokens=False,
)['input_ids'].squeeze(0)
def __getitem__(self, idx):
sample = self.samples[idx]
prompt = sample['prompt']
chosen = sample['chosen']
rejected = sample['rejected']
prompt_ids = self._tokenize(prompt)
chosen_ids = self._tokenize(prompt + chosen)
rejected_ids = self._tokenize(prompt + rejected)
prompt_len = len(prompt_ids)
return {
'chosen_ids': chosen_ids,
'rejected_ids': rejected_ids,
'prompt_len': prompt_len,
}
def collate_fn(batch, pad_id=0):
"""Pad sequences to same length within a batch."""
chosen_ids = [x['chosen_ids'] for x in batch]
rejected_ids = [x['rejected_ids'] for x in batch]
prompt_lens = [x['prompt_len'] for x in batch]
def pad(seqs):
max_len = max(len(s) for s in seqs)
padded = torch.zeros(len(seqs), max_len, dtype=torch.long)
mask = torch.zeros(len(seqs), max_len, dtype=torch.bool)
for i, s in enumerate(seqs):
padded[i, :len(s)] = s
mask[i, :len(s)] = True
return padded, mask
chosen_ids, chosen_mask = pad(chosen_ids)
rejected_ids, rejected_mask = pad(rejected_ids)
prompt_lens = torch.tensor(prompt_lens, dtype=torch.long)
return chosen_ids, chosen_mask, rejected_ids, rejected_mask, prompt_lens
# ── DPO loss ─────────────────────────────────────────────────────────────────
def compute_log_probs(model, input_ids, attention_mask, prompt_len):
"""
Compute sum of log-probs over response tokens only (not prompt).
Args:
model: language model
input_ids: (B, L)
attention_mask: (B, L)
prompt_len: (B,) — number of prompt tokens per sample
Returns:
log_probs: (B,) — sum of log-probs over response tokens
"""
with torch.no_grad() if not model.training else nullcontext():
# Model expects float attention_mask (0.0/1.0), not bool
attn_mask_float = attention_mask.float()
outputs = model(input_ids=input_ids, attention_mask=attn_mask_float)
logits = outputs.logits # (B, L, V)
# Shift: logits[i] predicts input_ids[i+1]
shift_logits = logits[:, :-1, :].contiguous() # (B, L-1, V)
shift_labels = input_ids[:, 1:].contiguous() # (B, L-1)
shift_mask = attention_mask[:, 1:].contiguous() # (B, L-1)
log_probs_all = F.log_softmax(shift_logits, dim=-1) # (B, L-1, V)
# Gather log-prob of the actual next token
token_log_probs = log_probs_all.gather(
2, shift_labels.unsqueeze(-1)
).squeeze(-1) # (B, L-1)
# Mask: only response tokens (after prompt)
response_mask = shift_mask.clone()
for i in range(len(prompt_len)):
# prompt_len[i]-1 because we shifted by 1
response_start = max(0, prompt_len[i].item() - 1)
response_mask[i, :response_start] = 0
# Sum log-probs over response tokens
token_log_probs = token_log_probs * response_mask.float()
log_probs = token_log_probs.sum(dim=-1) # (B,)
return log_probs
def dpo_loss(policy_chosen_logp, policy_rejected_logp,
ref_chosen_logp, ref_rejected_logp, beta: float = 0.1):
"""
DPO loss from Rafailov et al. (2023).
L = -E[log σ(β * ((log π(yw|x) - log πref(yw|x)) - (log π(yl|x) - log πref(yl|x))))]
"""
chosen_rewards = beta * (policy_chosen_logp - ref_chosen_logp)
rejected_rewards = beta * (policy_rejected_logp - ref_rejected_logp)
loss = -F.logsigmoid(chosen_rewards - rejected_rewards).mean()
# Preference accuracy: fraction of samples where chosen > rejected
with torch.no_grad():
pref_acc = (chosen_rewards > rejected_rewards).float().mean()
return loss, chosen_rewards.mean().detach(), rejected_rewards.mean().detach(), pref_acc
# ── Training ──────────────────────────────────────────────────────────────────
def train_epoch(epoch, loader, iters, policy_model, ref_model, optimizer,
scaler, autocast_ctx, args, wandb=None):
policy_model.train()
start_time = time.time()
total_loss = 0.0
total_acc = 0.0
for step, (chosen_ids, chosen_mask, rejected_ids, rejected_mask, prompt_lens) in \
enumerate(loader, start=1):
chosen_ids = chosen_ids.to(args.device)
chosen_mask = chosen_mask.to(args.device)
rejected_ids = rejected_ids.to(args.device)
rejected_mask = rejected_mask.to(args.device)
prompt_lens = prompt_lens.to(args.device)
lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
for pg in optimizer.param_groups:
pg['lr'] = lr
with autocast_ctx:
# Policy log-probs (with gradient)
policy_chosen_logp = compute_log_probs(
policy_model, chosen_ids, chosen_mask, prompt_lens)
policy_rejected_logp = compute_log_probs(
policy_model, rejected_ids, rejected_mask, prompt_lens)
# Reference log-probs (no gradient)
with torch.no_grad():
ref_chosen_logp = compute_log_probs(
ref_model, chosen_ids, chosen_mask, prompt_lens)
ref_rejected_logp = compute_log_probs(
ref_model, rejected_ids, rejected_mask, prompt_lens)
loss, chosen_r, rejected_r, pref_acc = dpo_loss(
policy_chosen_logp, policy_rejected_logp,
ref_chosen_logp, ref_rejected_logp,
beta=args.beta,
)
loss = loss / args.accumulation_steps
scaler.scale(loss).backward()
if step % args.accumulation_steps == 0:
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(policy_model.parameters(), args.grad_clip)
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad(set_to_none=True)
total_loss += loss.item() * args.accumulation_steps
total_acc += pref_acc.item()
if step % args.log_interval == 0 or step == iters:
elapsed = time.time() - start_time
avg_loss = total_loss / step
avg_acc = total_acc / step
Logger(
f'Epoch:[{epoch+1}/{args.epochs}]({step}/{iters}) '
f'loss:{avg_loss:.4f} pref_acc:{avg_acc:.3f} '
f'chosen_r:{chosen_r:.3f} rejected_r:{rejected_r:.3f} '
f'lr:{lr:.2e} time:{elapsed/60:.1f}min'
)
if wandb:
wandb.log({
'loss': avg_loss,
'preference_accuracy': avg_acc,
'chosen_reward': chosen_r.item(),
'rejected_reward': rejected_r.item(),
'lr': lr,
'epoch': epoch + 1,
})
return total_loss / iters, total_acc / iters
# ── Main ──────────────────────────────────────────────────────────────────────
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='DriveMind-V DPO Training')
parser.add_argument('--data_path', type=str, default='dataset/driving_preferences.jsonl')
parser.add_argument('--out_dir', type=str, default='out')
parser.add_argument('--epochs', type=int, default=3)
parser.add_argument('--batch_size', type=int, default=4)
parser.add_argument('--learning_rate', type=float, default=5e-5)
parser.add_argument('--beta', type=float, default=0.1,
help='DPO temperature (higher = closer to reference)')
parser.add_argument('--max_length', type=int, default=256)
parser.add_argument('--device', type=str,
default='cuda:0' if torch.cuda.is_available() else 'cpu')
parser.add_argument('--dtype', type=str,
default='bfloat16' if torch.cuda.is_bf16_supported() else 'float16')
parser.add_argument('--grad_clip', type=float, default=1.0)
parser.add_argument('--accumulation_steps', type=int, default=1)
parser.add_argument('--log_interval', type=int, default=5)
parser.add_argument('--num_workers', type=int, default=0)
parser.add_argument('--hidden_size', type=int, default=512)
parser.add_argument('--num_hidden_layers', type=int, default=8)
parser.add_argument('--from_weight', type=str, default='sft_vlm',
help='Base checkpoint prefix (loads out/<from_weight>_<hidden_size>.pth)')
# LoRA
parser.add_argument('--use_lora', action='store_true')
parser.add_argument('--lora_r', type=int, default=16)
parser.add_argument('--lora_alpha', type=int, default=32)
parser.add_argument('--lora_dropout', type=float, default=0.05)
# WandB
parser.add_argument('--use_wandb', action='store_true')
args = parser.parse_args()
setup_seed(42)
# ── Tokenizer ──
tokenizer = AutoTokenizer.from_pretrained('./model')
# ── Model config ──
vlm_config = VLMConfig(
hidden_size=args.hidden_size,
num_hidden_layers=args.num_hidden_layers,
)
# ── Load policy model ──
weight_path = f'{args.out_dir}/{args.from_weight}_{vlm_config.hidden_size}.pth'
Logger(f'Loading base model from {weight_path}')
policy_model = MiniMindVLM(vlm_config,
vision_model_path='./model/vision_model/clip-vit-base-patch16')
if os.path.exists(weight_path):
weights = torch.load(weight_path, map_location=args.device)
policy_model.load_state_dict(weights, strict=False)
Logger(f'Loaded weights from {weight_path}')
else:
Logger(f'WARNING: {weight_path} not found — using random init')
policy_model = policy_model.to(args.device)
# ── Frozen reference model (deep copy before LoRA) ──
Logger('Creating frozen reference model...')
ref_model = copy.deepcopy(policy_model)
ref_model.eval()
for p in ref_model.parameters():
p.requires_grad = False
Logger('Reference model frozen.')
# ── Apply LoRA to policy model ──
if args.use_lora:
Logger('\nApplying LoRA to policy model...')
lora_config = get_lora_config(
r=args.lora_r,
lora_alpha=args.lora_alpha,
lora_dropout=args.lora_dropout,
)
policy_model = apply_lora(policy_model, lora_config, verbose=True)
# ── Dataset & loader ──
dataset = DrivingPreferenceDataset(args.data_path, tokenizer, args.max_length)
loader = DataLoader(
dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
collate_fn=collate_fn,
pin_memory=(args.device != 'cpu'),
)
# ── Optimizer & AMP ──
optimizer = optim.AdamW(
filter(lambda p: p.requires_grad, policy_model.parameters()),
lr=args.learning_rate,
weight_decay=0.01,
)
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
if args.dtype == 'bfloat16':
autocast_ctx = torch.cuda.amp.autocast(dtype=torch.bfloat16)
elif args.dtype == 'float16':
autocast_ctx = torch.cuda.amp.autocast(dtype=torch.float16)
else:
autocast_ctx = nullcontext()
# ── WandB ──
wandb = None
if args.use_wandb:
import wandb as wb
wb.init(project='drivemind-v', name=f'dpo_lora_r{args.lora_r}_beta{args.beta}',
config=vars(args))
wandb = wb
# ── Train ──
iters = len(loader)
Logger(f'\nStarting DPO training: {args.epochs} epochs x {iters} steps')
Logger(f'Dataset: {len(dataset)} pairs | Batch: {args.batch_size} | Beta: {args.beta}')
Logger(f'LoRA: {args.use_lora} | Device: {args.device} | dtype: {args.dtype}\n')
for epoch in range(args.epochs):
avg_loss, avg_acc = train_epoch(
epoch, loader, iters,
policy_model, ref_model,
optimizer, scaler, autocast_ctx,
args, wandb,
)
Logger(f'Epoch {epoch+1} complete — avg_loss: {avg_loss:.4f} avg_pref_acc: {avg_acc:.3f}')
# ── Save ──
os.makedirs(args.out_dir, exist_ok=True)
lora_tag = '_lora' if args.use_lora else ''
save_path = f'{args.out_dir}/dpo_driving{lora_tag}_{vlm_config.hidden_size}'
if args.use_lora:
policy_model.save_pretrained(save_path + '_adapters')
Logger(f'LoRA adapters saved -> {save_path}_adapters')
else:
raw = getattr(policy_model, '_orig_mod', policy_model)
state_dict = {k: v.half().cpu() for k, v in raw.state_dict().items()
if not k.startswith('vision_encoder.')}
ckp = save_path + '.pth'
torch.save(state_dict, ckp)
Logger(f'Model saved -> {ckp}')
if wandb:
wandb.finish()
Logger('\nDPO training complete!')