diff --git a/dataset/driving_preferences.jsonl b/dataset/driving_preferences.jsonl new file mode 100644 index 0000000..aac0b14 --- /dev/null +++ b/dataset/driving_preferences.jsonl @@ -0,0 +1,56 @@ +{"category": "following_distance", "prompt": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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": "[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."} diff --git a/scripts/generate_driving_preferences.py b/scripts/generate_driving_preferences.py new file mode 100644 index 0000000..ca31aa1 --- /dev/null +++ b/scripts/generate_driving_preferences.py @@ -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"[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() diff --git a/trainer/train_dpo_driving.py b/trainer/train_dpo_driving.py new file mode 100644 index 0000000..8869e21 --- /dev/null +++ b/trainer/train_dpo_driving.py @@ -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/_.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!')