From c5b67f8c042d76423ace8bf3ee017b85a0377077 Mon Sep 17 00:00:00 2001 From: Bernhard Jaeger Date: Wed, 1 Jul 2026 19:03:04 +0200 Subject: [PATCH 01/35] add scripts and notes --- bernhard_configs.txt | 287 ++++++++++++++++++++++++ bernhard_notes.txt | 199 ++++++++++++++++ scripts/0_build_puffer.sh | 23 ++ scripts/0_eval_model_slurm.sh | 23 ++ scripts/0_train_model_slurm.sh | 23 ++ scripts/cluster_configs/bernhard_1.yaml | 7 + scripts/cluster_configs/bernhard_8.yaml | 7 + 7 files changed, 569 insertions(+) create mode 100644 bernhard_configs.txt create mode 100644 bernhard_notes.txt create mode 100644 scripts/0_build_puffer.sh create mode 100644 scripts/0_eval_model_slurm.sh create mode 100644 scripts/0_train_model_slurm.sh create mode 100644 scripts/cluster_configs/bernhard_1.yaml create mode 100644 scripts/cluster_configs/bernhard_8.yaml diff --git a/bernhard_configs.txt b/bernhard_configs.txt new file mode 100644 index 0000000000..71167db3fd --- /dev/null +++ b/bernhard_configs.txt @@ -0,0 +1,287 @@ +puffer train puffer_drive --vec.num-workers 1 --vec.num-envs 1 --env.num-agents 32 --train.minibatch-size 256 + +puffer train puffer_drive --train.learning_rate 0.001 + +.venv/bin/puffer train puffer_drive --wandb --vec.num-envs 16 --train.precision bfloat16 --train.compile 1 + +.venv/bin/puffer train puffer_drive --vec.num-envs 16 + +.venv/bin/puffer train puffer_drive --vec.num-envs 16 --train.precision bfloat16 --train.compile 1 --env.use-map-cache 1 + +uv run convert --py123d_path /home/jaeger/ordnung/data/nuPlan/123D --output /home/jaeger/ordnung/data/nuPlan/PufferDrive --preset nuplan + + +python scripts/submit_cluster.py --save_dir /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments \ + --prefix bernhard \ + --compute_config /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/scripts/cluster_configs/bernhard_1.yaml \ + --program_config /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/scripts/cluster_configs/train_base.yaml \ + --account bernhard --partition kyutai --time 2880 \ + --args \ + wandb=1 \ + vec.num-envs=24 + +python scripts/submit_cluster.py --save_dir /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments \ + --prefix bernhard \ + --compute_config /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/scripts/cluster_configs/bernhard_8.yaml \ + --program_config /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/scripts/cluster_configs/train_base.yaml \ + --account bernhard --partition kyutai --time 5760 \ + --args \ + wandb=1 \ + vec.num-envs=24 \ + train.total-timesteps=80000000000 + +python scripts/submit_cluster.py --save_dir /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments \ + --prefix bernhard \ + --compute_config /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/scripts/cluster_configs/bernhard_8.yaml \ + --program_config /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/scripts/cluster_configs/train_base.yaml \ + --account bernhard --partition kyutai --time 5760 \ + --args \ +run-name=k_004_longrun_300B \ +wandb=1 \ +train.total-timesteps=300000000000 \ +train.max-minibatch-size=131072 \ +train.minibatch-size=131072 \ +vec.num-envs=24 \ +eval.validation-replay.env.map-dir=/lustre/scwpod02/client/kyutai/kesai/data/nuPlan/PufferDrive \ +eval.behaviors-full-dir.env.map-dir=/lustre/scwpod02/client/kyutai/kesai/data/nuPlan/PufferDrive \ +env.dt=0.3 \ +env.max-agents-per-env=120 \ +env.num-agents=4096 \ +env.obs-boundary-stride=1 \ +env.obs-dropout-boundary=0.3 \ +env.obs-dropout-lane=0.2 \ +env.obs-lane-stride=1 \ +env.obs-norm-goal-offset-m=200.0 \ +env.obs-norm-veh-length-m=10.0 \ +env.obs-norm-veh-width-m=5.0 \ +env.obs-norm-xy-offset-m=150.0 \ +env.obs-range-partner-m=150.0 \ +env.obs-range-road-behind-m=40.0 \ +env.obs-range-road-front-m=150.0 \ +env.obs-range-road-side-m=50.0 \ +env.obs-range-traffic-control-m=150.0 \ +env.obs-slots-boundary-n=40 \ +env.obs-slots-lane-n=60 \ +env.obs-slots-partners-n=12 \ +env.resample-frequency=256000 \ +env.reward-goal=0.5 \ +env.reward-lane-center=0.005 \ +env.scenario-length=2560 \ +env.use-map-cache=1 \ +policy.backbone-hidden-size=1024 \ +policy.backbone-num-layers=3 \ +policy.ego-input-size=128 \ +policy.lane-input-size=128 \ +policy.partner-input-size=128 \ +policy.shared-network=0 \ +policy.traffic-control-input-size=128 \ +train.checkpoint-interval=50 \ +train.compile=1 \ +train.normalize-rewards=0 \ +train.precision=bfloat16 \ +train.seed=4 \ +train.update-epochs=3 + + \ +eval.validation-defaults.interval=1 + +env.traffic-control-scope=0 \ +policyboundary-input-size=128 \ +render-obs=0 \ + +pufferl.py: error: unrecognized arguments: --env.traffic-control-scope 0 --policyboundary-input-size 128 --render-obs 0 true + +Usage=puffer [--load-model-path LOAD_MODEL_PATH] [--load-id LOAD_ID] [--render-mode {auto,human,ansi,rgb_array,raylib,None}] [--video-path VIDEO_PATH] + [--num_scenarios NUM_SCENARIOS] [--render RENDER] [--agent_index [AGENT_INDEX ...]] [--save-frames SAVE_FRAMES] [--gif-path GIF_PATH] + [--fps FPS] [--max-runs MAX_RUNS] [--wandb] [--wandb-project WANDB_PROJECT] [--wandb-group WANDB_GROUP] [--run-name RUN_NAME] [--neptune] + [--neptune-name NEPTUNE_NAME] [--neptune-project NEPTUNE_PROJECT] [--tb] [--local-rank LOCAL_RANK] [--tag TAG] + [--eval_simulation EVAL_SIMULATION] [--package PACKAGE] [--env-name ENV_NAME] [--policy-name POLICY_NAME] [--rnn-name RNN_NAME] + [--max-suggestion-cost MAX_SUGGESTION_COST] [--vec.backend VEC.BACKEND] [--vec.num-envs VEC.NUM_ENVS] [--vec.num-workers VEC.NUM_WORKERS] + [--vec.batch-size VEC.BATCH_SIZE] [--vec.zero-copy VEC.ZERO_COPY] [--vec.seed VEC.SEED] [--env.simulation-mode ENV.SIMULATION_MODE] + [--env.num-agents ENV.NUM_AGENTS] [--env.min-agents-per-env ENV.MIN_AGENTS_PER_ENV] [--env.max-agents-per-env ENV.MAX_AGENTS_PER_ENV] + [--env.action-type ENV.ACTION_TYPE] [--env.dynamics-model ENV.DYNAMICS_MODEL] [--env.dt ENV.DT] + [--env.spawn-initial-speed ENV.SPAWN_INITIAL_SPEED] [--env.collision-behavior ENV.COLLISION_BEHAVIOR] + [--env.offroad-behavior ENV.OFFROAD_BEHAVIOR] [--env.traffic-light-behavior ENV.TRAFFIC_LIGHT_BEHAVIOR] + [--env.use-map-cache ENV.USE_MAP_CACHE] [--env.scenario-length ENV.SCENARIO_LENGTH] [--env.resample-frequency ENV.RESAMPLE_FREQUENCY] + [--env.termination-mode ENV.TERMINATION_MODE] [--env.inactive-agent-threshold ENV.INACTIVE_AGENT_THRESHOLD] + [--env.init-step ENV.INIT_STEP] [--env.control-mode ENV.CONTROL_MODE] [--env.sdc-controller ENV.SDC_CONTROLLER] + [--env.non-sdc-controller ENV.NON_SDC_CONTROLLER] [--env.non-vehicle-controller ENV.NON_VEHICLE_CONTROLLER] + [--env.init-mode ENV.INIT_MODE] [--env.compute-eval-metrics ENV.COMPUTE_EVAL_METRICS] [--env.target-type ENV.TARGET_TYPE] + [--env.goal-on-lane ENV.GOAL_ON_LANE] [--env.goal-radius ENV.GOAL_RADIUS] [--env.goal-speed ENV.GOAL_SPEED] + [--env.num-target-waypoints ENV.NUM_TARGET_WAYPOINTS] [--env.min-waypoint-spacing ENV.MIN_WAYPOINT_SPACING] + [--env.max-waypoint-spacing ENV.MAX_WAYPOINT_SPACING] [--env.reward-conditioning ENV.REWARD_CONDITIONING] + [--env.reward-randomization ENV.REWARD_RANDOMIZATION] [--env.reward-goal ENV.REWARD_GOAL] [--env.reward-collision ENV.REWARD_COLLISION] + [--env.reward-offroad ENV.REWARD_OFFROAD] [--env.reward-stop-line ENV.REWARD_STOP_LINE] [--env.reward-comfort ENV.REWARD_COMFORT] + [--env.reward-lane-align ENV.REWARD_LANE_ALIGN] [--env.reward-vel-align ENV.REWARD_VEL_ALIGN] + [--env.reward-lane-center ENV.REWARD_LANE_CENTER] [--env.reward-center-bias ENV.REWARD_CENTER_BIAS] + [--env.reward-velocity ENV.REWARD_VELOCITY] [--env.reward-reverse ENV.REWARD_REVERSE] [--env.reward-timestep ENV.REWARD_TIMESTEP] + [--env.reward-overspeed ENV.REWARD_OVERSPEED] [--env.reward-ade ENV.REWARD_ADE] [--env.map-dir ENV.MAP_DIR] [--env.num-maps ENV.NUM_MAPS] + [--env.obs-slots-lane-n ENV.OBS_SLOTS_LANE_N] [--env.obs-slots-boundary-n ENV.OBS_SLOTS_BOUNDARY_N] + [--env.obs-slots-partners-n ENV.OBS_SLOTS_PARTNERS_N] [--env.obs-slots-traffic-controls-n ENV.OBS_SLOTS_TRAFFIC_CONTROLS_N] + [--env.obs-dropout-lane ENV.OBS_DROPOUT_LANE] [--env.obs-dropout-boundary ENV.OBS_DROPOUT_BOUNDARY] + [--env.obs-norm-goal-offset-m ENV.OBS_NORM_GOAL_OFFSET_M] [--env.obs-norm-xy-offset-m ENV.OBS_NORM_XY_OFFSET_M] + [--env.obs-norm-veh-length-m ENV.OBS_NORM_VEH_LENGTH_M] [--env.obs-norm-veh-width-m ENV.OBS_NORM_VEH_WIDTH_M] + [--env.obs-norm-road-seg-length-m ENV.OBS_NORM_ROAD_SEG_LENGTH_M] [--env.obs-norm-road-seg-width-m ENV.OBS_NORM_ROAD_SEG_WIDTH_M] + [--env.obs-range-road-front-m ENV.OBS_RANGE_ROAD_FRONT_M] [--env.obs-range-road-behind-m ENV.OBS_RANGE_ROAD_BEHIND_M] + [--env.obs-range-road-side-m ENV.OBS_RANGE_ROAD_SIDE_M] [--env.obs-range-partner-m ENV.OBS_RANGE_PARTNER_M] + [--env.obs-range-traffic-control-m ENV.OBS_RANGE_TRAFFIC_CONTROL_M] [--env.partner-blindness-prob ENV.PARTNER_BLINDNESS_PROB] + [--env.partner-blindness-trigger-prob ENV.PARTNER_BLINDNESS_TRIGGER_PROB] [--env.phantom-braking-prob ENV.PHANTOM_BRAKING_PROB] + [--env.phantom-braking-trigger-prob ENV.PHANTOM_BRAKING_TRIGGER_PROB] [--env.phantom-braking-duration ENV.PHANTOM_BRAKING_DURATION] + [--policy.input-size POLICY.INPUT_SIZE] [--policy.encoder-gigaflow POLICY.ENCODER_GIGAFLOW] [--policy.dropout POLICY.DROPOUT] + [--policy.backbone-hidden-size POLICY.BACKBONE_HIDDEN_SIZE] [--policy.backbone-num-layers POLICY.BACKBONE_NUM_LAYERS] + [--policy.actor-hidden-size POLICY.ACTOR_HIDDEN_SIZE] [--policy.actor-num-layers POLICY.ACTOR_NUM_LAYERS] + [--policy.critic-hidden-size POLICY.CRITIC_HIDDEN_SIZE] [--policy.critic-num-layers POLICY.CRITIC_NUM_LAYERS] + [--policy.split-network POLICY.SPLIT_NETWORK] [--rnn.input-size RNN.INPUT_SIZE] [--rnn.hidden-size RNN.HIDDEN_SIZE] + [--train.name TRAIN.NAME] [--train.project TRAIN.PROJECT] [--train.seed TRAIN.SEED] + [--train.torch-deterministic TRAIN.TORCH_DETERMINISTIC] [--train.cpu-offload TRAIN.CPU_OFFLOAD] [--train.device TRAIN.DEVICE] + [--train.optimizer TRAIN.OPTIMIZER] [--train.anneal-lr TRAIN.ANNEAL_LR] [--train.amp TRAIN.AMP] [--train.precision TRAIN.PRECISION] + [--train.total-timesteps TRAIN.TOTAL_TIMESTEPS] [--train.learning-rate TRAIN.LEARNING_RATE] [--train.gamma TRAIN.GAMMA] + [--train.gae-lambda TRAIN.GAE_LAMBDA] [--train.update-epochs TRAIN.UPDATE_EPOCHS] [--train.clip-coef TRAIN.CLIP_COEF] + [--train.vf-coef TRAIN.VF_COEF] [--train.vf-clip-coef TRAIN.VF_CLIP_COEF] [--train.max-grad-norm TRAIN.MAX_GRAD_NORM] + [--train.normalize-rewards TRAIN.NORMALIZE_REWARDS] [--train.ent-coef TRAIN.ENT_COEF] [--train.adam-beta1 TRAIN.ADAM_BETA1] + [--train.adam-beta2 TRAIN.ADAM_BETA2] [--train.adam-eps TRAIN.ADAM_EPS] [--train.data-dir TRAIN.DATA_DIR] + [--train.resume-state-path TRAIN.RESUME_STATE_PATH] [--train.checkpoint-interval TRAIN.CHECKPOINT_INTERVAL] + [--train.batch-size TRAIN.BATCH_SIZE] [--train.minibatch-size TRAIN.MINIBATCH_SIZE] [--train.max-minibatch-size TRAIN.MAX_MINIBATCH_SIZE] + [--train.bptt-horizon TRAIN.BPTT_HORIZON] [--train.compile TRAIN.COMPILE] [--train.compile-mode TRAIN.COMPILE_MODE] + [--train.compile-fullgraph TRAIN.COMPILE_FULLGRAPH] [--train.vtrace-rho-clip TRAIN.VTRACE_RHO_CLIP] + [--train.vtrace-c-clip TRAIN.VTRACE_C_CLIP] [--train.adv-sampling-prio-alpha TRAIN.ADV_SAMPLING_PRIO_ALPHA] + [--train.adv-sampling-prio-beta0 TRAIN.ADV_SAMPLING_PRIO_BETA0] [--train.ppo-granularity TRAIN.PPO_GRANULARITY] + [--train.adv-filter-ewma-beta TRAIN.ADV_FILTER_EWMA_BETA] [--train.adv-filter-threshold-scale TRAIN.ADV_FILTER_THRESHOLD_SCALE] + [--train.render TRAIN.RENDER] [--train.render-interval TRAIN.RENDER_INTERVAL] [--train.obs-only TRAIN.OBS_ONLY] + [--train.show-grid TRAIN.SHOW_GRID] [--train.show-lasers TRAIN.SHOW_LASERS] [--train.show-human-logs TRAIN.SHOW_HUMAN_LOGS] + [--train.render-map TRAIN.RENDER_MAP] [--sweep.method SWEEP.METHOD] [--sweep.metric SWEEP.METRIC] [--sweep.goal SWEEP.GOAL] + [--sweep.downsample SWEEP.DOWNSAMPLE] [--mine.output-dir MINE.OUTPUT_DIR] [--mine.num-episodes MINE.NUM_EPISODES] + [--mine.score-threshold MINE.SCORE_THRESHOLD] [--mine.render MINE.RENDER] + [--eval.validation-defaults.enabled EVAL.VALIDATION_DEFAULTS.ENABLED] + [--eval.validation-defaults.interval EVAL.VALIDATION_DEFAULTS.INTERVAL] [--eval.validation-defaults.mode EVAL.VALIDATION_DEFAULTS.MODE] + [--eval.validation-defaults.clean EVAL.VALIDATION_DEFAULTS.CLEAN] + [--eval.validation-defaults.env.eval-mode EVAL.VALIDATION_DEFAULTS.ENV.EVAL_MODE] + [--eval.validation-defaults.env.collision-behavior EVAL.VALIDATION_DEFAULTS.ENV.COLLISION_BEHAVIOR] + [--eval.validation-defaults.env.offroad-behavior EVAL.VALIDATION_DEFAULTS.ENV.OFFROAD_BEHAVIOR] + [--eval.validation-defaults.env.traffic-light-behavior EVAL.VALIDATION_DEFAULTS.ENV.TRAFFIC_LIGHT_BEHAVIOR] + [--eval.validation-defaults.env.reward-randomization EVAL.VALIDATION_DEFAULTS.ENV.REWARD_RANDOMIZATION] + [--eval.validation-defaults.env.termination-mode EVAL.VALIDATION_DEFAULTS.ENV.TERMINATION_MODE] + [--eval.validation-defaults.env.num-agents EVAL.VALIDATION_DEFAULTS.ENV.NUM_AGENTS] + [--eval.validation-defaults.env.target-type EVAL.VALIDATION_DEFAULTS.ENV.TARGET_TYPE] + [--eval.validation-defaults.env.goal-speed EVAL.VALIDATION_DEFAULTS.ENV.GOAL_SPEED] + [--eval.validation-defaults.env.reward-collision EVAL.VALIDATION_DEFAULTS.ENV.REWARD_COLLISION] + [--eval.validation-defaults.env.reward-offroad EVAL.VALIDATION_DEFAULTS.ENV.REWARD_OFFROAD] + [--eval.validation-defaults.env.reward-stop-line EVAL.VALIDATION_DEFAULTS.ENV.REWARD_STOP_LINE] + [--eval.validation-defaults.env.reward-goal EVAL.VALIDATION_DEFAULTS.ENV.REWARD_GOAL] + [--eval.validation-defaults.env.reward-overspeed EVAL.VALIDATION_DEFAULTS.ENV.REWARD_OVERSPEED] + [--eval.validation-defaults.env.reward-comfort EVAL.VALIDATION_DEFAULTS.ENV.REWARD_COMFORT] + [--eval.validation-defaults.env.reward-velocity EVAL.VALIDATION_DEFAULTS.ENV.REWARD_VELOCITY] + [--eval.validation-defaults.env.reward-lane-align EVAL.VALIDATION_DEFAULTS.ENV.REWARD_LANE_ALIGN] + [--eval.validation-defaults.env.reward-lane-center EVAL.VALIDATION_DEFAULTS.ENV.REWARD_LANE_CENTER] + [--eval.validation-defaults.env.reward-timestep EVAL.VALIDATION_DEFAULTS.ENV.REWARD_TIMESTEP] + [--eval.validation-defaults.env.reward-reverse EVAL.VALIDATION_DEFAULTS.ENV.REWARD_REVERSE] + [--eval.validation-defaults.env.reward-ade EVAL.VALIDATION_DEFAULTS.ENV.REWARD_ADE] + [--eval.validation-defaults.env.obs-dropout-lane EVAL.VALIDATION_DEFAULTS.ENV.OBS_DROPOUT_LANE] + [--eval.validation-defaults.env.obs-dropout-boundary EVAL.VALIDATION_DEFAULTS.ENV.OBS_DROPOUT_BOUNDARY] + [--eval.validation-defaults.env.obs-slots-lane-n EVAL.VALIDATION_DEFAULTS.ENV.OBS_SLOTS_LANE_N] + [--eval.validation-defaults.env.obs-slots-boundary-n EVAL.VALIDATION_DEFAULTS.ENV.OBS_SLOTS_BOUNDARY_N] + [--eval.validation-defaults.eval.num-scenarios EVAL.VALIDATION_DEFAULTS.EVAL.NUM_SCENARIOS] + [--eval.validation-defaults.eval.export-episode-csv EVAL.VALIDATION_DEFAULTS.EVAL.EXPORT_EPISODE_CSV] + [--eval.validation-defaults.eval.verify-coverage EVAL.VALIDATION_DEFAULTS.EVAL.VERIFY_COVERAGE] + [--eval.validation-replay.inherits EVAL.VALIDATION_REPLAY.INHERITS] [--eval.validation-replay.type EVAL.VALIDATION_REPLAY.TYPE] + [--eval.validation-replay.enabled EVAL.VALIDATION_REPLAY.ENABLED] [--eval.validation-replay.render EVAL.VALIDATION_REPLAY.RENDER] + [--eval.validation-replay.render-backend EVAL.VALIDATION_REPLAY.RENDER_BACKEND] + [--eval.validation-replay.env.simulation-mode EVAL.VALIDATION_REPLAY.ENV.SIMULATION_MODE] + [--eval.validation-replay.env.control-mode EVAL.VALIDATION_REPLAY.ENV.CONTROL_MODE] + [--eval.validation-replay.env.map-dir EVAL.VALIDATION_REPLAY.ENV.MAP_DIR] + [--eval.validation-replay.env.num-maps EVAL.VALIDATION_REPLAY.ENV.NUM_MAPS] + [--eval.validation-replay.env.max-agents-per-env EVAL.VALIDATION_REPLAY.ENV.MAX_AGENTS_PER_ENV] + [--eval.validation-replay.env.scenario-length EVAL.VALIDATION_REPLAY.ENV.SCENARIO_LENGTH] + [--eval.validation-replay.env.resample-frequency EVAL.VALIDATION_REPLAY.ENV.RESAMPLE_FREQUENCY] + [--eval.validation-replay.eval.render-num-scenarios EVAL.VALIDATION_REPLAY.EVAL.RENDER_NUM_SCENARIOS] + [--eval.validation-replay.eval.render-max-steps EVAL.VALIDATION_REPLAY.EVAL.RENDER_MAX_STEPS] + [--eval.validation-gigaflow.inherits EVAL.VALIDATION_GIGAFLOW.INHERITS] [--eval.validation-gigaflow.type EVAL.VALIDATION_GIGAFLOW.TYPE] + [--eval.validation-gigaflow.enabled EVAL.VALIDATION_GIGAFLOW.ENABLED] [--eval.validation-gigaflow.render EVAL.VALIDATION_GIGAFLOW.RENDER] + [--eval.validation-gigaflow.render-backend EVAL.VALIDATION_GIGAFLOW.RENDER_BACKEND] + [--eval.validation-gigaflow.render-views EVAL.VALIDATION_GIGAFLOW.RENDER_VIEWS] + [--eval.validation-gigaflow.env.simulation-mode EVAL.VALIDATION_GIGAFLOW.ENV.SIMULATION_MODE] + [--eval.validation-gigaflow.env.map-dir EVAL.VALIDATION_GIGAFLOW.ENV.MAP_DIR] + [--eval.validation-gigaflow.env.num-maps EVAL.VALIDATION_GIGAFLOW.ENV.NUM_MAPS] + [--eval.validation-gigaflow.env.num-agents EVAL.VALIDATION_GIGAFLOW.ENV.NUM_AGENTS] + [--eval.validation-gigaflow.env.min-agents-per-env EVAL.VALIDATION_GIGAFLOW.ENV.MIN_AGENTS_PER_ENV] + [--eval.validation-gigaflow.env.max-agents-per-env EVAL.VALIDATION_GIGAFLOW.ENV.MAX_AGENTS_PER_ENV] + [--eval.validation-gigaflow.env.scenario-length EVAL.VALIDATION_GIGAFLOW.ENV.SCENARIO_LENGTH] + [--eval.validation-gigaflow.env.resample-frequency EVAL.VALIDATION_GIGAFLOW.ENV.RESAMPLE_FREQUENCY] + [--eval.validation-gigaflow.eval.render-num-scenarios EVAL.VALIDATION_GIGAFLOW.EVAL.RENDER_NUM_SCENARIOS] + [--eval.validation-gigaflow.eval.render-max-steps EVAL.VALIDATION_GIGAFLOW.EVAL.RENDER_MAX_STEPS] + [--eval.dnf-triage.inherits EVAL.DNF_TRIAGE.INHERITS] [--eval.dnf-triage.type EVAL.DNF_TRIAGE.TYPE] + [--eval.dnf-triage.enabled EVAL.DNF_TRIAGE.ENABLED] [--eval.dnf-triage.render EVAL.DNF_TRIAGE.RENDER] + [--eval.dnf-triage.render-backend EVAL.DNF_TRIAGE.RENDER_BACKEND] + [--eval.dnf-triage.env.simulation-mode EVAL.DNF_TRIAGE.ENV.SIMULATION_MODE] [--eval.dnf-triage.env.map-dir EVAL.DNF_TRIAGE.ENV.MAP_DIR] + [--eval.dnf-triage.env.num-maps EVAL.DNF_TRIAGE.ENV.NUM_MAPS] + [--eval.dnf-triage.env.min-agents-per-env EVAL.DNF_TRIAGE.ENV.MIN_AGENTS_PER_ENV] + [--eval.dnf-triage.env.max-agents-per-env EVAL.DNF_TRIAGE.ENV.MAX_AGENTS_PER_ENV] + [--eval.dnf-triage.env.scenario-length EVAL.DNF_TRIAGE.ENV.SCENARIO_LENGTH] + [--eval.dnf-triage.env.resample-frequency EVAL.DNF_TRIAGE.ENV.RESAMPLE_FREQUENCY] + [--eval.dnf-triage.eval.num-scenarios EVAL.DNF_TRIAGE.EVAL.NUM_SCENARIOS] + [--eval.dnf-triage.eval.render-num-scenarios EVAL.DNF_TRIAGE.EVAL.RENDER_NUM_SCENARIOS] + [--eval.dnf-triage.eval.render-max-steps EVAL.DNF_TRIAGE.EVAL.RENDER_MAX_STEPS] + [--eval.behaviors-defaults.enabled EVAL.BEHAVIORS_DEFAULTS.ENABLED] [--eval.behaviors-defaults.interval EVAL.BEHAVIORS_DEFAULTS.INTERVAL] + [--eval.behaviors-defaults.mode EVAL.BEHAVIORS_DEFAULTS.MODE] [--eval.behaviors-defaults.clean EVAL.BEHAVIORS_DEFAULTS.CLEAN] + [--eval.behaviors-defaults.render EVAL.BEHAVIORS_DEFAULTS.RENDER] + [--eval.behaviors-defaults.render-views EVAL.BEHAVIORS_DEFAULTS.RENDER_VIEWS] + [--eval.behaviors-defaults.env.simulation-mode EVAL.BEHAVIORS_DEFAULTS.ENV.SIMULATION_MODE] + [--eval.behaviors-defaults.env.control-mode EVAL.BEHAVIORS_DEFAULTS.ENV.CONTROL_MODE] + [--eval.behaviors-defaults.env.init-mode EVAL.BEHAVIORS_DEFAULTS.ENV.INIT_MODE] + [--eval.behaviors-defaults.env.scenario-length EVAL.BEHAVIORS_DEFAULTS.ENV.SCENARIO_LENGTH] + [--eval.behaviors-defaults.env.obs-slots-partners-n EVAL.BEHAVIORS_DEFAULTS.ENV.OBS_SLOTS_PARTNERS_N] + [--eval.behaviors-defaults.eval.num-scenarios EVAL.BEHAVIORS_DEFAULTS.EVAL.NUM_SCENARIOS] + [--eval.behaviors-defaults.eval.render-num-scenarios EVAL.BEHAVIORS_DEFAULTS.EVAL.RENDER_NUM_SCENARIOS] + [--eval.behaviors-defaults.eval.render-max-steps EVAL.BEHAVIORS_DEFAULTS.EVAL.RENDER_MAX_STEPS] + [--eval.behaviors-full-dir.inherits EVAL.BEHAVIORS_FULL_DIR.INHERITS] [--eval.behaviors-full-dir.type EVAL.BEHAVIORS_FULL_DIR.TYPE] + [--eval.behaviors-full-dir.enabled EVAL.BEHAVIORS_FULL_DIR.ENABLED] + [--eval.behaviors-full-dir.env.map-dir EVAL.BEHAVIORS_FULL_DIR.ENV.MAP_DIR] + [--eval.behaviors-hard-stop.inherits EVAL.BEHAVIORS_HARD_STOP.INHERITS] [--eval.behaviors-hard-stop.type EVAL.BEHAVIORS_HARD_STOP.TYPE] + [--eval.behaviors-hard-stop.enabled EVAL.BEHAVIORS_HARD_STOP.ENABLED] + [--eval.behaviors-hard-stop.env.map-dir EVAL.BEHAVIORS_HARD_STOP.ENV.MAP_DIR] + [--eval.behaviors-highway-straight.inherits EVAL.BEHAVIORS_HIGHWAY_STRAIGHT.INHERITS] + [--eval.behaviors-highway-straight.type EVAL.BEHAVIORS_HIGHWAY_STRAIGHT.TYPE] + [--eval.behaviors-highway-straight.enabled EVAL.BEHAVIORS_HIGHWAY_STRAIGHT.ENABLED] + [--eval.behaviors-highway-straight.env.map-dir EVAL.BEHAVIORS_HIGHWAY_STRAIGHT.ENV.MAP_DIR] + [--eval.behaviors-lane-change.inherits EVAL.BEHAVIORS_LANE_CHANGE.INHERITS] + [--eval.behaviors-lane-change.type EVAL.BEHAVIORS_LANE_CHANGE.TYPE] + [--eval.behaviors-lane-change.enabled EVAL.BEHAVIORS_LANE_CHANGE.ENABLED] + [--eval.behaviors-lane-change.env.map-dir EVAL.BEHAVIORS_LANE_CHANGE.ENV.MAP_DIR] + [--eval.behaviors-merge.inherits EVAL.BEHAVIORS_MERGE.INHERITS] [--eval.behaviors-merge.type EVAL.BEHAVIORS_MERGE.TYPE] + [--eval.behaviors-merge.enabled EVAL.BEHAVIORS_MERGE.ENABLED] [--eval.behaviors-merge.env.map-dir EVAL.BEHAVIORS_MERGE.ENV.MAP_DIR] + [--eval.behaviors-parked-cars.inherits EVAL.BEHAVIORS_PARKED_CARS.INHERITS] + [--eval.behaviors-parked-cars.type EVAL.BEHAVIORS_PARKED_CARS.TYPE] + [--eval.behaviors-parked-cars.enabled EVAL.BEHAVIORS_PARKED_CARS.ENABLED] + [--eval.behaviors-parked-cars.env.map-dir EVAL.BEHAVIORS_PARKED_CARS.ENV.MAP_DIR] + [--eval.behaviors-roundabout.inherits EVAL.BEHAVIORS_ROUNDABOUT.INHERITS] + [--eval.behaviors-roundabout.type EVAL.BEHAVIORS_ROUNDABOUT.TYPE] [--eval.behaviors-roundabout.enabled EVAL.BEHAVIORS_ROUNDABOUT.ENABLED] + [--eval.behaviors-roundabout.env.map-dir EVAL.BEHAVIORS_ROUNDABOUT.ENV.MAP_DIR] + [--eval.behaviors-stopped-traffic.inherits EVAL.BEHAVIORS_STOPPED_TRAFFIC.INHERITS] + [--eval.behaviors-stopped-traffic.type EVAL.BEHAVIORS_STOPPED_TRAFFIC.TYPE] + [--eval.behaviors-stopped-traffic.enabled EVAL.BEHAVIORS_STOPPED_TRAFFIC.ENABLED] + [--eval.behaviors-stopped-traffic.env.map-dir EVAL.BEHAVIORS_STOPPED_TRAFFIC.ENV.MAP_DIR] + [--eval.behaviors-traffic-light-green.inherits EVAL.BEHAVIORS_TRAFFIC_LIGHT_GREEN.INHERITS] + [--eval.behaviors-traffic-light-green.type EVAL.BEHAVIORS_TRAFFIC_LIGHT_GREEN.TYPE] + [--eval.behaviors-traffic-light-green.enabled EVAL.BEHAVIORS_TRAFFIC_LIGHT_GREEN.ENABLED] + [--eval.behaviors-traffic-light-green.env.map-dir EVAL.BEHAVIORS_TRAFFIC_LIGHT_GREEN.ENV.MAP_DIR] + [--eval.behaviors-traffic-light-stop.inherits EVAL.BEHAVIORS_TRAFFIC_LIGHT_STOP.INHERITS] + [--eval.behaviors-traffic-light-stop.type EVAL.BEHAVIORS_TRAFFIC_LIGHT_STOP.TYPE] + [--eval.behaviors-traffic-light-stop.enabled EVAL.BEHAVIORS_TRAFFIC_LIGHT_STOP.ENABLED] + [--eval.behaviors-traffic-light-stop.env.map-dir EVAL.BEHAVIORS_TRAFFIC_LIGHT_STOP.ENV.MAP_DIR] + [--eval.behaviors-unprotected-left.inherits EVAL.BEHAVIORS_UNPROTECTED_LEFT.INHERITS] + [--eval.behaviors-unprotected-left.type EVAL.BEHAVIORS_UNPROTECTED_LEFT.TYPE] + [--eval.behaviors-unprotected-left.enabled EVAL.BEHAVIORS_UNPROTECTED_LEFT.ENABLED] + [--eval.behaviors-unprotected-left.env.map-dir EVAL.BEHAVIORS_UNPROTECTED_LEFT.ENV.MAP_DIR] + [--eval.behaviors-unprotected-right.inherits EVAL.BEHAVIORS_UNPROTECTED_RIGHT.INHERITS] + [--eval.behaviors-unprotected-right.type EVAL.BEHAVIORS_UNPROTECTED_RIGHT.TYPE] + [--eval.behaviors-unprotected-right.enabled EVAL.BEHAVIORS_UNPROTECTED_RIGHT.ENABLED] + [--eval.behaviors-unprotected-right.env.map-dir EVAL.BEHAVIORS_UNPROTECTED_RIGHT.ENV.MAP_DIR] [--eval.wosac.type EVAL.WOSAC.TYPE] + [--eval.wosac.enabled EVAL.WOSAC.ENABLED] [--eval.wosac.interval EVAL.WOSAC.INTERVAL] [--eval.wosac.mode EVAL.WOSAC.MODE] + [--eval.wosac.clean EVAL.WOSAC.CLEAN] [--eval.wosac.render EVAL.WOSAC.RENDER] [--eval.wosac.env.control-mode EVAL.WOSAC.ENV.CONTROL_MODE] + [--eval.wosac.env.init-mode EVAL.WOSAC.ENV.INIT_MODE] [--eval.wosac.env.init-step EVAL.WOSAC.ENV.INIT_STEP] + [--eval.wosac.env.goal-radius EVAL.WOSAC.ENV.GOAL_RADIUS] [--eval.wosac.eval.wosac-num-rollouts EVAL.WOSAC.EVAL.WOSAC_NUM_ROLLOUTS] + [--eval.wosac.eval.wosac-num-agents EVAL.WOSAC.EVAL.WOSAC_NUM_AGENTS] + [--eval.wosac.eval.wosac-sanity-check EVAL.WOSAC.EVAL.WOSAC_SANITY_CHECK] + [--eval.wosac.eval.wosac-aggregate-results EVAL.WOSAC.EVAL.WOSAC_AGGREGATE_RESULTS] + [--controlled-exp.train.learning-rate.values CONTROLLED_EXP.TRAIN.LEARNING_RATE.VALUES] + [--controlled-exp.train.ent-coef.values CONTROLLED_EXP.TRAIN.ENT_COEF.VALUES] [-h] \ No newline at end of file diff --git a/bernhard_notes.txt b/bernhard_notes.txt new file mode 100644 index 0000000000..6a2ebcc07c --- /dev/null +++ b/bernhard_notes.txt @@ -0,0 +1,199 @@ +git merge origin/3.0 +git push origin bernhard + + +./scripts/setup_container.sh --scratch-dir /lustre/scwpod02/client/kyutai/kesai/bernhard/puffer_scratch create-overlay +export OVERLAY_PATH=/lustre/scwpod02/client/kyutai/kesai/bernhard/puffer_scratch/overlay-15GB-500K.ext3 +export IMAGE_PATH=/lustre/scwpod02/client/kyutai/kesai/bernhard/puffer_scratch/cuda12.8.1-cudnn9.8.0-ubuntu24.04.2.sif +export OVERLAY_TEMPLATE=/lustre/scwpod02/client/kyutai/kesai/bernhard/puffer_scratch/overlay-15GB-500K.ext3.gz +export VENV_PATH=/lustre/scwpod02/client/kyutai/kesai/bernhard/puffer_scratch/venvs/pufferdrive +export CONTAINER_PYTHON=/home/bernhard/miniconda3/envs/puffer/bin/python3.14 + +"train", + "puffer_drive", + "--wandb", + "--vec.num-envs", "16", + +SPS 313.4K +SPS 341.4K + + + + File "/lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/bernhard_train_base_wandb1_run-name003_longrun_total-timesteps300000000000_num-envs24_max-waypoint-spacing40.0_num-agents2048_obs-norm-ve_067a078/code_v1/pufferlib/ocean/benchmark/manager.py", line 123, in maybe_run + res = self._run_one( + File "/lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/bernhard_train_base_wandb1_run-name003_longrun_total-timesteps300000000000_num-envs24_max-waypoint-spacing40.0_num-agents2048_obs-norm-ve_067a078/code_v1/pufferlib/ocean/benchmark/manager.py", line 148, in _run_one + res = self._run_inline(ev, policy=policy, env_name=env_name, global_step=global_step, epoch=epoch) + File "/lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/bernhard_train_base_wandb1_run-name003_longrun_total-timesteps300000000000_num-envs24_max-waypoint-spacing40.0_num-agents2048_obs-norm-ve_067a078/code_v1/pufferlib/ocean/benchmark/manager.py", line 189, in _run_inline + res = ev.rollout(vecenv, policy, args) + File "/lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/bernhard_train_base_wandb1_run-name003_longrun_total-timesteps300000000000_num-envs24_max-waypoint-spacing40.0_num-agents2048_obs-norm-ve_067a078/code_v1/pufferlib/ocean/benchmark/evaluators/base.py", line 113, in rollout + metrics = self._run_rollout_loop(vecenv, policy, args) + File "/lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/bernhard_train_base_wandb1_run-name003_longrun_total-timesteps300000000000_num-envs24_max-waypoint-spacing40.0_num-agents2048_obs-norm-ve_067a078/code_v1/pufferlib/ocean/benchmark/evaluators/base.py", line 156, in _run_rollout_loop + logits, _ = policy.forward_eval(ob_t, state) + + + + + + slide_idx 687 +traffic_control_dim 28 +self.obs_valid_count_features 4 +self.obs_slots_traffic_controls_n 4 +self.traffic_control_features_count 7 +traffic_control_observations.shape torch.Size([393216, 28]) +self.traffic_control_continuous_features + 1 6 + + + +slide_idx 687 +traffic_control_dim 28 +self.obs_valid_count_features 4 +self.obs_slots_traffic_controls_n 4 +self.traffic_control_features_count 7 +traffic_control_observations.shape torch.Size([1024, 28]) +self.traffic_control_continuous_features + 1 6 +[eval.validation_replay] wrote 250 per-episode rows to episode_metrics/validation_replay_epoch1_step6291456.csv +slide_idx 687 +traffic_control_dim 28 +self.obs_valid_count_features 4 +self.obs_slots_traffic_controls_n 4 +self.traffic_control_features_count 7 +traffic_control_observations.shape torch.Size([1024, 28]) +self.traffic_control_continuous_features + 1 6 + +Expected max limit (num_classes): 5 +Min index encountered: -0.999995231628418 +Max index encountered: 0.9999845027923584 +Expected max limit (num_classes): 5 +Min index encountered: -1.0 +Max index encountered: 1.0 +Expected max limit (num_classes): 5 +[EvalManager] Evaluator 'validation_gigaflow' raised at epoch 1; continuing with the rest: + + + + + +TODO read puffer advantage. +Why is my puffer not compiled with CUDA advantage? +ADVANTAGE_CUDA + +what is adv_filter_ewma_beta? Is that the advantage filtering of GIGAFLOW + +lauch.json: + +{ + "version"="0.2.0", + "configurations"=[ + { + "name"="PufferDrive Train", + "type"="python", + "request"="launch", + "program"="${workspaceFolder}/.venv/bin/puffer", + "args"=[ + "train", + "puffer_drive", + "--vec.num-envs", "16", + "--env.num-agents", "1024", + "--train.minibatch-size", "131072", + "--train.max-minibatch-size", "131072", + ], + "console"="integratedTerminal", + "justMyCode"=false + } + ] +} + + +# Run a named evaluator on a checkpoint (config from [eval.]) +puffer eval puffer_drive --evaluator validation_gigaflow \ + --load-model-path experiments/puffer_drive_xxxx/models/model_puffer_drive_000500.pt + +# Ad-hoc=pick by simulation + override scale from the CLI +puffer eval puffer_drive --eval_simulation replay \ + --load-model-path experiments/puffer_drive_xxxx/models/model_puffer_drive_000500.pt \ + --num_scenarios 250 --render 1 + +# Render the agent's observations (interactive HTML) +puffer eval puffer_drive --eval_simulation gigaflow \ + --load-model-path experiments/puffer_drive_xxxx/models/model_puffer_drive_000500.pt \ + --num_scenarios 10 --render 1 --render-backend obs_html + + +Eval nodes CARLA: +Traffic lights are only red for an extremely short amount of time. Policies almost never have to wait (although they do). +Models seem to deviate from lane center occasionally. +Policies can drive backward. + + +╭───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮ +│ PufferLib 3.0 🐡 CPU=95.7% GPU=0.0% DRAM=44.0% VRAM=0.0% │ +│ │ +│ Summary Value Evaluate 1m 16s 48% Losses Value │ +│ Env puffer_drive Forward 25s 15% filter_threshold 0.047 │ +│ Params 1.0M Env 28s 18% ema_max 4.726 │ +│ Steps 60.8M Copy 22s 13% masked_fraction 0.144 │ +│ SPS 343.7K Misc 0s 0% kept_fraction 0.764 │ +│ Epoch 29 Train 1m 25s 51% filtered_fraction 0.236 │ +│ Uptime 3m 12s Forward 36s 21% policy_loss -0.001 │ +│ Remaining 8h 1m 55s Learn 47s 28% value_loss 0.065 │ +│ Copy 0s 0% entropy 1.937 │ +│ Misc 0s 0% old_approx_kl 0.006 │ +│ approx_kl 0.006 │ +│ clipfrac 0.065 │ +│ explained_variance 0.926 │ +│ │ +│ User Stats Value User Stats Value │ +│ obs/max 4.000 obs/min -2.516 │ +│ obs/mean 0.180 n 64.000 │ +│ offroad_rate 0.179 episode_length 113.655 │ +│ collision_rate 0.178 episode_return 0.672 │ +│ red_light_violation_rate 0.054 comfort_violation_count 0.001 │ +│ velocity_progress_sum 0.617 num_goals_reached 0.095 │ +│ lane_center_rate 0.678 dnf_rate 0.324 │ +│ score 0.000 avg_speed_per_agent 4.854 │ +│ avg_distance_per_infraction 145.226 reward_components/collision -0.373 │ +│ reward_components/offroad -0.269 reward_components/red_light -0.054 │ +│ reward_components/goal 0.669 reward_components/lane_align -0.003 │ +│ reward_components/lane_center -0.009 reward_components/comfort -0.009 │ +│ reward_components/velocity 0.722 reward_components/timestep -0.000 │ +│ reward_components/reverse -0.000 reward_components/overspeed -0.002 │ +│ reward_components/ade 0.000 │ +╰───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ + + +With Expo and PBO + +╭───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮ +│ PufferLib 3.0 🐡 CPU=93.5% GPU=0.0% DRAM=43.0% VRAM=0.0% │ +│ │ +│ Summary Value Evaluate 1m 23s 48% Losses Value │ +│ Env puffer_drive Forward 29s 17% filter_threshold 0.053 │ +│ Params 1.0M Env 30s 17% ema_max 5.300 │ +│ Steps 65.0M Copy 22s 13% masked_fraction 0.151 │ +│ SPS 363.7K Misc 0s 0% kept_fraction 0.746 │ +│ Epoch 31 Train 1m 40s 51% filtered_fraction 0.254 │ +│ Uptime 3m 16s Forward 42s 22% policy_loss 0.000 │ +│ Remaining 7h 35m 14s Learn 56s 28% value_loss 0.079 │ +│ Copy 0s 0% entropy 1.930 │ +│ Misc 0s 0% old_approx_kl 0.006 │ +│ approx_kl 0.006 │ +│ clipfrac 0.076 │ +│ explained_variance 0.925 │ +│ │ +│ User Stats Value User Stats Value │ +│ obs/max 4.000 obs/min -2.434 │ +│ obs/mean 0.179 n 64.000 │ +│ offroad_rate 0.150 episode_length 126.518 │ +│ collision_rate 0.197 episode_return 0.837 │ +│ red_light_violation_rate 0.064 comfort_violation_count 0.002 │ +│ velocity_progress_sum 0.618 num_goals_reached 0.118 │ +│ lane_center_rate 0.657 dnf_rate 0.312 │ +│ score 0.000 avg_speed_per_agent 4.699 │ +│ avg_distance_per_infraction 155.931 reward_components/collision -0.403 │ +│ reward_components/offroad -0.225 reward_components/red_light -0.064 │ +│ reward_components/goal 0.756 reward_components/lane_align -0.003 │ +│ reward_components/lane_center -0.011 reward_components/comfort -0.011 │ +│ reward_components/velocity 0.802 reward_components/timestep -0.000 │ +│ reward_components/reverse -0.000 reward_components/overspeed -0.003 │ +│ reward_components/ade 0.000 │ +╰───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ \ No newline at end of file diff --git a/scripts/0_build_puffer.sh b/scripts/0_build_puffer.sh new file mode 100644 index 0000000000..92020d5e25 --- /dev/null +++ b/scripts/0_build_puffer.sh @@ -0,0 +1,23 @@ +#!/bin/bash +#SBATCH --job-name=train_puffer +#SBATCH --ntasks=1 +#SBATCH --nodes=1 +#SBATCH --time=1-00:00 +#SBATCH --gres=gpu:1 +#SBATCH --mem=64G +#SBATCH --cpus-per-task=24 +#SBATCH --output=/lustre/scwpod02/client/kyutai/kesai/bernhard/carla_closed_loop/results/logs/eval_server_%a_%A.out +#SBATCH --error=/lustre/scwpod02/client/kyutai/kesai/bernhard/carla_closed_loop/results/logs/eval_server_%a_%A.err +#SBATCH --partition=kyutai + +# print info about current job +echo "START TIME: $(date)" +start=`date +%s` + +source .venv/bin/activate +python setup.py build_ext --inplace --force + +end=`date +%s` +runtime=$((end-start)) +echo "END TIME: $(date)" +echo "Runtime: ${runtime}" \ No newline at end of file diff --git a/scripts/0_eval_model_slurm.sh b/scripts/0_eval_model_slurm.sh new file mode 100644 index 0000000000..c8d49b7692 --- /dev/null +++ b/scripts/0_eval_model_slurm.sh @@ -0,0 +1,23 @@ +#!/bin/bash +#SBATCH --job-name=train_puffer +#SBATCH --ntasks=1 +#SBATCH --nodes=1 +#SBATCH --time=1-00:00 +#SBATCH --gres=gpu:1 +#SBATCH --mem=256G +#SBATCH --cpus-per-task=28 +#SBATCH --output=/lustre/scwpod02/client/kyutai/kesai/bernhard/carla_closed_loop/results/logs/eval_server_%a_%A.out +#SBATCH --error=/lustre/scwpod02/client/kyutai/kesai/bernhard/carla_closed_loop/results/logs/eval_server_%a_%A.err +#SBATCH --partition=kyutai + +# print info about current job +echo "START TIME: $(date)" +start=`date +%s` + +source .venv/bin/activate +.venv/bin/puffer eval puffer_drive --eval_simulation replay --render 1 --render-backend obs_html --load-model-path /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/bernhard_train_base_wandb1_run-namek_004_longrun_300B_total-timesteps300000000000_max-minibatch-size131072_minibatch-size131072_num-envs2_91bb92d/puffer_drive_j671z22p/best_models/best_trainer_state_002981.pt + +end=`date +%s` +runtime=$((end-start)) +echo "END TIME: $(date)" +echo "Runtime: ${runtime}" \ No newline at end of file diff --git a/scripts/0_train_model_slurm.sh b/scripts/0_train_model_slurm.sh new file mode 100644 index 0000000000..7c0c89613b --- /dev/null +++ b/scripts/0_train_model_slurm.sh @@ -0,0 +1,23 @@ +#!/bin/bash +#SBATCH --job-name=train_puffer +#SBATCH --ntasks=1 +#SBATCH --nodes=1 +#SBATCH --time=1-00:00 +#SBATCH --gres=gpu:1 +#SBATCH --mem=128G +#SBATCH --cpus-per-task=28 +#SBATCH --output=/lustre/scwpod02/client/kyutai/kesai/bernhard/carla_closed_loop/results/logs/eval_server_%a_%A.out +#SBATCH --error=/lustre/scwpod02/client/kyutai/kesai/bernhard/carla_closed_loop/results/logs/eval_server_%a_%A.err +#SBATCH --partition=kyutai + +# print info about current job +echo "START TIME: $(date)" +start=`date +%s` + +source .venv/bin/activate +.venv/bin/puffer train puffer_drive --wandb --vec.num-envs 28 + +end=`date +%s` +runtime=$((end-start)) +echo "END TIME: $(date)" +echo "Runtime: ${runtime}" \ No newline at end of file diff --git a/scripts/cluster_configs/bernhard_1.yaml b/scripts/cluster_configs/bernhard_1.yaml new file mode 100644 index 0000000000..90b8798bd2 --- /dev/null +++ b/scripts/cluster_configs/bernhard_1.yaml @@ -0,0 +1,7 @@ +# Kyutai cluster compute configuration +account: bernhard # Set your account/allocation +nodes: 1 +gpus: 1 +cpus: 28 +mem: 128gb +time: 2880 # minutes \ No newline at end of file diff --git a/scripts/cluster_configs/bernhard_8.yaml b/scripts/cluster_configs/bernhard_8.yaml new file mode 100644 index 0000000000..5bd0b7f4dc --- /dev/null +++ b/scripts/cluster_configs/bernhard_8.yaml @@ -0,0 +1,7 @@ +# Kyutai cluster compute configuration +account: bernhard # Set your account/allocation +nodes: 1 +gpus: 8 +cpus: 224 +mem: 1024gb +time: 5760 # minutes \ No newline at end of file From 10d9e554012fe1c1bb3e224f78b034de17a5d6f5 Mon Sep 17 00:00:00 2001 From: Bernhard Jaeger Date: Fri, 3 Jul 2026 12:53:01 +0200 Subject: [PATCH 02/35] Change nuPlan folders.eps --- .gitignore | 2 + bernhard_configs.txt | 57 ++++++++++++- pufferlib/config/ocean/drive.ini | 136 +++++++++++++++---------------- 3 files changed, 126 insertions(+), 69 deletions(-) diff --git a/.gitignore b/.gitignore index baa6caa7bf..4c374516bb 100644 --- a/.gitignore +++ b/.gitignore @@ -218,3 +218,5 @@ docs/_build/ # Claude config .claude/ CLAUDE.local.md + +episode_metrics \ No newline at end of file diff --git a/bernhard_configs.txt b/bernhard_configs.txt index 71167db3fd..eed75cb470 100644 --- a/bernhard_configs.txt +++ b/bernhard_configs.txt @@ -6,7 +6,7 @@ puffer train puffer_drive --train.learning_rate 0.001 .venv/bin/puffer train puffer_drive --vec.num-envs 16 -.venv/bin/puffer train puffer_drive --vec.num-envs 16 --train.precision bfloat16 --train.compile 1 --env.use-map-cache 1 +.venv/bin/puffer train puffer_drive --vec.num-envs 16 --train.precision bfloat16 --train.compile 1 --env.use-map-cache 1 --env.reward-randomization=1 --env.reward-conditioning=1 --train.total-timesteps=30000000000 --wandb --run-name b_005_local_30B_reward_random uv run convert --py123d_path /home/jaeger/ordnung/data/nuPlan/123D --output /home/jaeger/ordnung/data/nuPlan/PufferDrive --preset nuplan @@ -82,6 +82,61 @@ train.precision=bfloat16 \ train.seed=4 \ train.update-epochs=3 + +python scripts/submit_cluster.py --save_dir /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments \ + --prefix bernhard \ + --compute_config /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/scripts/cluster_configs/bernhard_8.yaml \ + --program_config /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/scripts/cluster_configs/train_base.yaml \ + --account bernhard --partition kyutai --time 5760 \ + --args \ +run-name=k_004_longrun_300B \ +wandb=1 \ +train.total-timesteps=300000000000 \ +train.max-minibatch-size=131072 \ +train.minibatch-size=131072 \ +vec.num-envs=24 \ +eval.validation-replay.env.map-dir=/lustre/scwpod02/client/kyutai/kesai/data/nuPlan/PufferDrive \ +eval.behaviors-full-dir.env.map-dir=/lustre/scwpod02/client/kyutai/kesai/data/nuPlan/PufferDrive \ +env.dt=0.3 \ +env.max-agents-per-env=120 \ +env.num-agents=4096 \ +env.obs-boundary-stride=1 \ +env.obs-dropout-boundary=0.3 \ +env.obs-dropout-lane=0.2 \ +env.obs-lane-stride=1 \ +env.obs-norm-goal-offset-m=200.0 \ +env.obs-norm-veh-length-m=10.0 \ +env.obs-norm-veh-width-m=5.0 \ +env.obs-norm-xy-offset-m=150.0 \ +env.obs-range-partner-m=150.0 \ +env.obs-range-road-behind-m=40.0 \ +env.obs-range-road-front-m=150.0 \ +env.obs-range-road-side-m=50.0 \ +env.obs-range-traffic-control-m=150.0 \ +env.obs-slots-boundary-n=40 \ +env.obs-slots-lane-n=60 \ +env.obs-slots-partners-n=12 \ +env.resample-frequency=256000 \ +env.reward-goal=0.5 \ +env.reward-lane-center=0.005 \ +env.scenario-length=2560 \ +env.reward-randomization=1 \ +env.reward-conditioning=1 \ +env.use-map-cache=1 \ +policy.backbone-hidden-size=1024 \ +policy.backbone-num-layers=3 \ +policy.ego-input-size=128 \ +policy.lane-input-size=128 \ +policy.partner-input-size=128 \ +policy.shared-network=0 \ +policy.traffic-control-input-size=128 \ +train.checkpoint-interval=50 \ +train.compile=1 \ +train.normalize-rewards=0 \ +train.precision=bfloat16 \ +train.seed=4 \ +train.update-epochs=3 + \ eval.validation-defaults.interval=1 diff --git a/pufferlib/config/ocean/drive.ini b/pufferlib/config/ocean/drive.ini index 5283573d38..7934f8266e 100644 --- a/pufferlib/config/ocean/drive.ini +++ b/pufferlib/config/ocean/drive.ini @@ -298,7 +298,7 @@ render = true render_backend = "triage_html" env.simulation_mode = "replay" env.control_mode = "control_sdc_only" -env.map_dir = "/scratch/ev2237/data/nuplan/nuplan_mini_train_bins" +env.map_dir = "/lustre/scwpod02/client/kyutai/kesai/data/nuPlan/PufferDrive" env.num_maps = 250 env.max_agents_per_env = 64 env.scenario_length = 200 @@ -367,73 +367,73 @@ eval.render_max_steps = 200 inherits = "behaviors_defaults" type = "behavior_class" enabled = true -env.map_dir = "/scratch/ev2237/data/nuplan/nuplan_mini_train_bins" - -[eval.behaviors_hard_stop] -inherits = "behaviors_defaults" -type = "behavior_class" -enabled = true -env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/hard_stop" - -[eval.behaviors_highway_straight] -inherits = "behaviors_defaults" -type = "behavior_class" -enabled = true -env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/highway_straight" - -[eval.behaviors_lane_change] -inherits = "behaviors_defaults" -type = "behavior_class" -enabled = true -env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/lane_change" - -[eval.behaviors_merge] -inherits = "behaviors_defaults" -type = "behavior_class" -enabled = true -env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/merge" - -[eval.behaviors_parked_cars] -inherits = "behaviors_defaults" -type = "behavior_class" -enabled = true -env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/parked_cars" - -[eval.behaviors_roundabout] -inherits = "behaviors_defaults" -type = "behavior_class" -enabled = true -env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/roundabout" - -[eval.behaviors_stopped_traffic] -inherits = "behaviors_defaults" -type = "behavior_class" -enabled = true -env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/stopped_traffic" - -[eval.behaviors_traffic_light_green] -inherits = "behaviors_defaults" -type = "behavior_class" -enabled = true -env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/traffic_light_green" - -[eval.behaviors_traffic_light_stop] -inherits = "behaviors_defaults" -type = "behavior_class" -enabled = true -env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/traffic_light_stop" - -[eval.behaviors_unprotected_left] -inherits = "behaviors_defaults" -type = "behavior_class" -enabled = true -env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/unprotected_left" - -[eval.behaviors_unprotected_right] -inherits = "behaviors_defaults" -type = "behavior_class" -enabled = true -env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/unprotected_right" +env.map_dir = "/lustre/scwpod02/client/kyutai/kesai/data/nuPlan/PufferDrive" + +; [eval.behaviors_hard_stop] +; inherits = "behaviors_defaults" +; type = "behavior_class" +; enabled = true +; env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/hard_stop" +; +; [eval.behaviors_highway_straight] +; inherits = "behaviors_defaults" +; type = "behavior_class" +; enabled = true +; env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/highway_straight" +; +; [eval.behaviors_lane_change] +; inherits = "behaviors_defaults" +; type = "behavior_class" +; enabled = true +; env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/lane_change" +; +; [eval.behaviors_merge] +; inherits = "behaviors_defaults" +; type = "behavior_class" +; enabled = true +; env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/merge" +; +; [eval.behaviors_parked_cars] +; inherits = "behaviors_defaults" +; type = "behavior_class" +; enabled = true +; env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/parked_cars" +; +; [eval.behaviors_roundabout] +; inherits = "behaviors_defaults" +; type = "behavior_class" +; enabled = true +; env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/roundabout" +; +; [eval.behaviors_stopped_traffic] +; inherits = "behaviors_defaults" +; type = "behavior_class" +; enabled = true +; env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/stopped_traffic" +; +; [eval.behaviors_traffic_light_green] +; inherits = "behaviors_defaults" +; type = "behavior_class" +; enabled = true +; env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/traffic_light_green" +; +; [eval.behaviors_traffic_light_stop] +; inherits = "behaviors_defaults" +; type = "behavior_class" +; enabled = true +; env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/traffic_light_stop" +; +; [eval.behaviors_unprotected_left] +; inherits = "behaviors_defaults" +; type = "behavior_class" +; enabled = true +; env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/unprotected_left" +; +; [eval.behaviors_unprotected_right] +; inherits = "behaviors_defaults" +; type = "behavior_class" +; enabled = true +; env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/unprotected_right" ; --------------------------------------------------------------------------- ; Optional: WOSAC realism eval. Off by default. From f41d573b493da88ab37606d965c00264cb0f3356 Mon Sep 17 00:00:00 2001 From: Bernhard Jaeger Date: Fri, 3 Jul 2026 13:32:55 +0200 Subject: [PATCH 03/35] add direct training script. --- bernhard_configs.txt | 58 -------------------------- scripts/0_train_model_slurm.sh | 75 +++++++++++++++++++++++++++------- 2 files changed, 61 insertions(+), 72 deletions(-) diff --git a/bernhard_configs.txt b/bernhard_configs.txt index eed75cb470..dfd5c81e5d 100644 --- a/bernhard_configs.txt +++ b/bernhard_configs.txt @@ -82,64 +82,6 @@ train.precision=bfloat16 \ train.seed=4 \ train.update-epochs=3 - -python scripts/submit_cluster.py --save_dir /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments \ - --prefix bernhard \ - --compute_config /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/scripts/cluster_configs/bernhard_8.yaml \ - --program_config /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/scripts/cluster_configs/train_base.yaml \ - --account bernhard --partition kyutai --time 5760 \ - --args \ -run-name=k_004_longrun_300B \ -wandb=1 \ -train.total-timesteps=300000000000 \ -train.max-minibatch-size=131072 \ -train.minibatch-size=131072 \ -vec.num-envs=24 \ -eval.validation-replay.env.map-dir=/lustre/scwpod02/client/kyutai/kesai/data/nuPlan/PufferDrive \ -eval.behaviors-full-dir.env.map-dir=/lustre/scwpod02/client/kyutai/kesai/data/nuPlan/PufferDrive \ -env.dt=0.3 \ -env.max-agents-per-env=120 \ -env.num-agents=4096 \ -env.obs-boundary-stride=1 \ -env.obs-dropout-boundary=0.3 \ -env.obs-dropout-lane=0.2 \ -env.obs-lane-stride=1 \ -env.obs-norm-goal-offset-m=200.0 \ -env.obs-norm-veh-length-m=10.0 \ -env.obs-norm-veh-width-m=5.0 \ -env.obs-norm-xy-offset-m=150.0 \ -env.obs-range-partner-m=150.0 \ -env.obs-range-road-behind-m=40.0 \ -env.obs-range-road-front-m=150.0 \ -env.obs-range-road-side-m=50.0 \ -env.obs-range-traffic-control-m=150.0 \ -env.obs-slots-boundary-n=40 \ -env.obs-slots-lane-n=60 \ -env.obs-slots-partners-n=12 \ -env.resample-frequency=256000 \ -env.reward-goal=0.5 \ -env.reward-lane-center=0.005 \ -env.scenario-length=2560 \ -env.reward-randomization=1 \ -env.reward-conditioning=1 \ -env.use-map-cache=1 \ -policy.backbone-hidden-size=1024 \ -policy.backbone-num-layers=3 \ -policy.ego-input-size=128 \ -policy.lane-input-size=128 \ -policy.partner-input-size=128 \ -policy.shared-network=0 \ -policy.traffic-control-input-size=128 \ -train.checkpoint-interval=50 \ -train.compile=1 \ -train.normalize-rewards=0 \ -train.precision=bfloat16 \ -train.seed=4 \ -train.update-epochs=3 - - \ -eval.validation-defaults.interval=1 - env.traffic-control-scope=0 \ policyboundary-input-size=128 \ render-obs=0 \ diff --git a/scripts/0_train_model_slurm.sh b/scripts/0_train_model_slurm.sh index 7c0c89613b..ee02b3acad 100644 --- a/scripts/0_train_model_slurm.sh +++ b/scripts/0_train_model_slurm.sh @@ -1,23 +1,70 @@ #!/bin/bash -#SBATCH --job-name=train_puffer -#SBATCH --ntasks=1 -#SBATCH --nodes=1 -#SBATCH --time=1-00:00 -#SBATCH --gres=gpu:1 -#SBATCH --mem=128G -#SBATCH --cpus-per-task=28 -#SBATCH --output=/lustre/scwpod02/client/kyutai/kesai/bernhard/carla_closed_loop/results/logs/eval_server_%a_%A.out -#SBATCH --error=/lustre/scwpod02/client/kyutai/kesai/bernhard/carla_closed_loop/results/logs/eval_server_%a_%A.err -#SBATCH --partition=kyutai +#SBATCH --job-name train_puffer +#SBATCH --ntasks 1 +#SBATCH --nodes 1 +#SBATCH --time 2-00:00 +#SBATCH --gres gpu:8 +#SBATCH --mem 1024G +#SBATCH --cpus-per-task 224 +#SBATCH --output /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/k_005_longrun_300B_%a_%A.out +#SBATCH --error /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/k_005_longrun_300B_%a_%A.err +#SBATCH --partition kyutai # print info about current job echo "START TIME: $(date)" -start=`date +%s` +start `date +%s` + +export RUN_NAME=k_005_longrun_300B +echo ${RUN_NAME} source .venv/bin/activate -.venv/bin/puffer train puffer_drive --wandb --vec.num-envs 28 +torchrun --standalone --nnodes=1 --nproc-per-node=8 --max_restarts=0 --start-method spawn -m pufferlib.pufferl train puffer_drive --wandb --vec.num-envs 28 --train.data-dir /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/${RUN_NAME} \ + --wandb --wandb-project pufferdrive --wandb-group cluster --train.checkpoint-interval 1000 --run-name ${RUN_NAME} \ + --train.total-timesteps 300000000000 \ + --train.max-minibatch-size 131072 \ + --train.minibatch-size 131072 \ + --vec.num-envs 24 \ + --eval.validation-replay.env.map-dir /lustre/scwpod02/client/kyutai/kesai/data/nuPlan/PufferDrive \ + --eval.behaviors-full-dir.env.map-dir /lustre/scwpod02/client/kyutai/kesai/data/nuPlan/PufferDrive \ + --env.dt 0.3 \ + --env.max-agents-per-env 120 \ + --env.num-agents 4096 \ + --env.obs-boundary-stride 1 \ + --env.obs-dropout-boundary 0.3 \ + --env.obs-dropout-lane 0.2 \ + --env.obs-lane-stride 1 \ + --env.obs-norm-goal-offset-m 200.0 \ + --env.obs-norm-veh-length-m 10.0 \ + --env.obs-norm-veh-width-m 5.0 \ + --env.obs-norm-xy-offset-m 150.0 \ + --env.obs-range-partner-m 150.0 \ + --env.obs-range-road-behind-m 40.0 \ + --env.obs-range-road-front-m 150.0 \ + --env.obs-range-road-side-m 50.0 \ + --env.obs-range-traffic-control-m 150.0 \ + --env.obs-slots-boundary-n 40 \ + --env.obs-slots-lane-n 60 \ + --env.obs-slots-partners-n 12 \ + --env.resample-frequency 256000 \ + --env.reward-goal 0.5 \ + --env.reward-lane-center 0.005 \ + --env.scenario-length 2560 \ + --env.use-map-cache 1 \ + --policy.backbone-hidden-size 1024 \ + --policy.backbone-num-layers 3 \ + --policy.ego-input-size 128 \ + --policy.lane-input-size 128 \ + --policy.partner-input-size 128 \ + --policy.shared-network 0 \ + --policy.traffic-control-input-size 128 \ + --train.checkpoint-interval 50 \ + --train.compile 1 \ + --train.normalize-rewards 0 \ + --train.precision bfloat16 \ + --train.seed 4 \ + --train.update-epochs 3 -end=`date +%s` -runtime=$((end-start)) +end `date +%s` +runtime $((end-start)) echo "END TIME: $(date)" echo "Runtime: ${runtime}" \ No newline at end of file From a233b2d070e0bf1e77839ecd85aa372218e57f45 Mon Sep 17 00:00:00 2001 From: Bernhard Jaeger Date: Fri, 3 Jul 2026 13:37:22 +0200 Subject: [PATCH 04/35] change folder --- scripts/0_train_model_slurm.sh | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/scripts/0_train_model_slurm.sh b/scripts/0_train_model_slurm.sh index ee02b3acad..ee90acbeec 100644 --- a/scripts/0_train_model_slurm.sh +++ b/scripts/0_train_model_slurm.sh @@ -6,13 +6,13 @@ #SBATCH --gres gpu:8 #SBATCH --mem 1024G #SBATCH --cpus-per-task 224 -#SBATCH --output /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/k_005_longrun_300B_%a_%A.out -#SBATCH --error /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/k_005_longrun_300B_%a_%A.err +#SBATCH --output /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/k_005_longrun_300B/log_%a_%A.out +#SBATCH --error /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/k_005_longrun_300B/log_%a_%A.err #SBATCH --partition kyutai # print info about current job echo "START TIME: $(date)" -start `date +%s` +start=$(date +%s) export RUN_NAME=k_005_longrun_300B echo ${RUN_NAME} @@ -64,7 +64,7 @@ torchrun --standalone --nnodes=1 --nproc-per-node=8 --max_restarts=0 --start-met --train.seed 4 \ --train.update-epochs 3 -end `date +%s` -runtime $((end-start)) +end=$(date +%s) +runtime=$((end-start)) echo "END TIME: $(date)" echo "Runtime: ${runtime}" \ No newline at end of file From 95fda3b19256b4440dcd6a4286e0159bb87f6755 Mon Sep 17 00:00:00 2001 From: Bernhard Jaeger Date: Fri, 3 Jul 2026 13:45:16 +0200 Subject: [PATCH 05/35] remove extra envs. --- scripts/0_train_model_slurm.sh | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/scripts/0_train_model_slurm.sh b/scripts/0_train_model_slurm.sh index ee90acbeec..6f0bdadb40 100644 --- a/scripts/0_train_model_slurm.sh +++ b/scripts/0_train_model_slurm.sh @@ -17,8 +17,15 @@ start=$(date +%s) export RUN_NAME=k_005_longrun_300B echo ${RUN_NAME} +# TODO could try to tune these. 1 Is probably best since Puffer parallelizes across all cores. +export NUMEXPR_NUM_THREADS=1 +export MKL_NUM_THREADS=1 +export VECLIB_MAXIMUM_THREADS=1 +export OPENBLAS_NUM_THREADS=1 +export OMP_NUM_THREADS=1 + source .venv/bin/activate -torchrun --standalone --nnodes=1 --nproc-per-node=8 --max_restarts=0 --start-method spawn -m pufferlib.pufferl train puffer_drive --wandb --vec.num-envs 28 --train.data-dir /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/${RUN_NAME} \ +torchrun --standalone --nnodes=1 --nproc-per-node=8 --max_restarts=0 --start-method spawn -m pufferlib.pufferl train puffer_drive --wandb --train.data-dir /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/${RUN_NAME} \ --wandb --wandb-project pufferdrive --wandb-group cluster --train.checkpoint-interval 1000 --run-name ${RUN_NAME} \ --train.total-timesteps 300000000000 \ --train.max-minibatch-size 131072 \ From 883ccc63c58259f737c5261a8e53b096ca397c02 Mon Sep 17 00:00:00 2001 From: Bernhard Jaeger Date: Fri, 3 Jul 2026 14:19:54 +0200 Subject: [PATCH 06/35] update params. --- scripts/0_train_model_slurm.sh | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/scripts/0_train_model_slurm.sh b/scripts/0_train_model_slurm.sh index 6f0bdadb40..728d8af86e 100644 --- a/scripts/0_train_model_slurm.sh +++ b/scripts/0_train_model_slurm.sh @@ -14,7 +14,7 @@ echo "START TIME: $(date)" start=$(date +%s) -export RUN_NAME=k_005_longrun_300B +export RUN_NAME=k_006_longrun_300B echo ${RUN_NAME} # TODO could try to tune these. 1 Is probably best since Puffer parallelizes across all cores. @@ -30,7 +30,8 @@ torchrun --standalone --nnodes=1 --nproc-per-node=8 --max_restarts=0 --start-met --train.total-timesteps 300000000000 \ --train.max-minibatch-size 131072 \ --train.minibatch-size 131072 \ - --vec.num-envs 24 \ + --vec.num-envs 48 \ + --vec.num-workers 24 \ --eval.validation-replay.env.map-dir /lustre/scwpod02/client/kyutai/kesai/data/nuPlan/PufferDrive \ --eval.behaviors-full-dir.env.map-dir /lustre/scwpod02/client/kyutai/kesai/data/nuPlan/PufferDrive \ --env.dt 0.3 \ @@ -62,14 +63,16 @@ torchrun --standalone --nnodes=1 --nproc-per-node=8 --max_restarts=0 --start-met --policy.ego-input-size 128 \ --policy.lane-input-size 128 \ --policy.partner-input-size 128 \ - --policy.shared-network 0 \ + --policy.shared-network false \ --policy.traffic-control-input-size 128 \ --train.checkpoint-interval 50 \ - --train.compile 1 \ - --train.normalize-rewards 0 \ + --train.compile true \ + --train.normalize-rewards false \ --train.precision bfloat16 \ --train.seed 4 \ - --train.update-epochs 3 + --train.update-epochs 3 \ + --policy.boundary-input-size 128 \ + --tb end=$(date +%s) runtime=$((end-start)) From ab98fc2902a9968950b90779d527799110b5dda3 Mon Sep 17 00:00:00 2001 From: Bernhard Jaeger Date: Fri, 3 Jul 2026 14:33:48 +0200 Subject: [PATCH 07/35] less batch size.- --- bernhard_notes.txt | 19 +++++++++++++++++++ scripts/0_train_model_slurm.sh | 4 ++-- 2 files changed, 21 insertions(+), 2 deletions(-) diff --git a/bernhard_notes.txt b/bernhard_notes.txt index 6a2ebcc07c..0e3ea21dde 100644 --- a/bernhard_notes.txt +++ b/bernhard_notes.txt @@ -2,6 +2,25 @@ git merge origin/3.0 git push origin bernhard +# Not in valentins config but in my config. +# non_sdc_controller: policy +# non_vehicle_controller: auto +# sdc_controller: policy +# goal_on_lane: true +# mine: +# num_episodes: 100 +# output_dir: '' +# render: 'true' +# score_threshold: -inf +# train: +# amp true + +# Not in my config but in valentins config: +# traffic_control_scope: 0 +# num_maps: 1 +# render_mode: matplotlib +# render_obs: 0 + ./scripts/setup_container.sh --scratch-dir /lustre/scwpod02/client/kyutai/kesai/bernhard/puffer_scratch create-overlay export OVERLAY_PATH=/lustre/scwpod02/client/kyutai/kesai/bernhard/puffer_scratch/overlay-15GB-500K.ext3 export IMAGE_PATH=/lustre/scwpod02/client/kyutai/kesai/bernhard/puffer_scratch/cuda12.8.1-cudnn9.8.0-ubuntu24.04.2.sif diff --git a/scripts/0_train_model_slurm.sh b/scripts/0_train_model_slurm.sh index 728d8af86e..db9b2f7512 100644 --- a/scripts/0_train_model_slurm.sh +++ b/scripts/0_train_model_slurm.sh @@ -28,8 +28,8 @@ source .venv/bin/activate torchrun --standalone --nnodes=1 --nproc-per-node=8 --max_restarts=0 --start-method spawn -m pufferlib.pufferl train puffer_drive --wandb --train.data-dir /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/${RUN_NAME} \ --wandb --wandb-project pufferdrive --wandb-group cluster --train.checkpoint-interval 1000 --run-name ${RUN_NAME} \ --train.total-timesteps 300000000000 \ - --train.max-minibatch-size 131072 \ - --train.minibatch-size 131072 \ + --train.max-minibatch-size 100000 \ + --train.minibatch-size 100000 \ --vec.num-envs 48 \ --vec.num-workers 24 \ --eval.validation-replay.env.map-dir /lustre/scwpod02/client/kyutai/kesai/data/nuPlan/PufferDrive \ From 4f5f979d4339c6f825a53a268be45a9cf0fa5b51 Mon Sep 17 00:00:00 2001 From: Bernhard Jaeger Date: Fri, 3 Jul 2026 15:04:48 +0200 Subject: [PATCH 08/35] update script. --- bernhard_configs.txt | 52 ++++++++++++++++++++++++++++++++++ scripts/0_train_model_slurm.sh | 10 +++---- 2 files changed, 57 insertions(+), 5 deletions(-) diff --git a/bernhard_configs.txt b/bernhard_configs.txt index dfd5c81e5d..525fa1f0cc 100644 --- a/bernhard_configs.txt +++ b/bernhard_configs.txt @@ -8,6 +8,58 @@ puffer train puffer_drive --train.learning_rate 0.001 .venv/bin/puffer train puffer_drive --vec.num-envs 16 --train.precision bfloat16 --train.compile 1 --env.use-map-cache 1 --env.reward-randomization=1 --env.reward-conditioning=1 --train.total-timesteps=30000000000 --wandb --run-name b_005_local_30B_reward_random + +torchrun --standalone --nnodes=1 --nproc-per-node=1 --max_restarts=0 --start-method spawn -m pufferlib.pufferl train puffer_drive --wandb --train.data-dir /home/jaeger/ordnung/internal/PufferDrive/experiments/debug_000 \ + --wandb --wandb-project pufferdrive --wandb-group cluster --train.checkpoint-interval 1000 --run-name debug_000 \ + --train.total-timesteps 30000000000 \ + --train.max-minibatch-size 65536 \ + --train.minibatch-size 65536 \ + --vec.num-envs 16 \ + --vec.num-workers 16 \ + --eval.validation-replay.env.map-dir /lustre/scwpod02/client/kyutai/kesai/data/nuPlan/PufferDrive \ + --eval.behaviors-full-dir.env.map-dir /lustre/scwpod02/client/kyutai/kesai/data/nuPlan/PufferDrive \ + --env.dt 0.3 \ + --env.max-agents-per-env 120 \ + --env.num-agents 4096 \ + --env.obs-boundary-stride 1 \ + --env.obs-dropout-boundary 0.3 \ + --env.obs-dropout-lane 0.2 \ + --env.obs-lane-stride 1 \ + --env.obs-norm-goal-offset-m 200.0 \ + --env.obs-norm-veh-length-m 10.0 \ + --env.obs-norm-veh-width-m 5.0 \ + --env.obs-norm-xy-offset-m 150.0 \ + --env.obs-range-partner-m 150.0 \ + --env.obs-range-road-behind-m 40.0 \ + --env.obs-range-road-front-m 150.0 \ + --env.obs-range-road-side-m 50.0 \ + --env.obs-range-traffic-control-m 150.0 \ + --env.obs-slots-boundary-n 40 \ + --env.obs-slots-lane-n 60 \ + --env.obs-slots-partners-n 12 \ + --env.resample-frequency 256000 \ + --env.reward-goal 0.5 \ + --env.reward-lane-center 0.005 \ + --env.scenario-length 2560 \ + --env.use-map-cache 1 \ + --policy.backbone-hidden-size 1024 \ + --policy.backbone-num-layers 3 \ + --policy.ego-input-size 128 \ + --policy.lane-input-size 128 \ + --policy.partner-input-size 128 \ + --policy.shared-network false \ + --policy.traffic-control-input-size 128 \ + --train.checkpoint-interval 50 \ + --train.compile true \ + --train.normalize-rewards false \ + --train.precision bfloat16 \ + --train.seed 4 \ + --train.update-epochs 3 \ + --policy.boundary-input-size 128 \ + --tb \ + --eval.validation-replay.enabled false \ + --eval.behaviors-full-dir.enabled false + uv run convert --py123d_path /home/jaeger/ordnung/data/nuPlan/123D --output /home/jaeger/ordnung/data/nuPlan/PufferDrive --preset nuplan diff --git a/scripts/0_train_model_slurm.sh b/scripts/0_train_model_slurm.sh index db9b2f7512..abf8ea34b5 100644 --- a/scripts/0_train_model_slurm.sh +++ b/scripts/0_train_model_slurm.sh @@ -6,8 +6,8 @@ #SBATCH --gres gpu:8 #SBATCH --mem 1024G #SBATCH --cpus-per-task 224 -#SBATCH --output /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/k_005_longrun_300B/log_%a_%A.out -#SBATCH --error /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/k_005_longrun_300B/log_%a_%A.err +#SBATCH --output /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/k_006_longrun_300B/log_%a_%A.out +#SBATCH --error /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/k_006_longrun_300B/log_%a_%A.err #SBATCH --partition kyutai # print info about current job @@ -28,9 +28,9 @@ source .venv/bin/activate torchrun --standalone --nnodes=1 --nproc-per-node=8 --max_restarts=0 --start-method spawn -m pufferlib.pufferl train puffer_drive --wandb --train.data-dir /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/${RUN_NAME} \ --wandb --wandb-project pufferdrive --wandb-group cluster --train.checkpoint-interval 1000 --run-name ${RUN_NAME} \ --train.total-timesteps 300000000000 \ - --train.max-minibatch-size 100000 \ - --train.minibatch-size 100000 \ - --vec.num-envs 48 \ + --train.max-minibatch-size 128000 \ + --train.minibatch-size 128000 \ + --vec.num-envs 24 \ --vec.num-workers 24 \ --eval.validation-replay.env.map-dir /lustre/scwpod02/client/kyutai/kesai/data/nuPlan/PufferDrive \ --eval.behaviors-full-dir.env.map-dir /lustre/scwpod02/client/kyutai/kesai/data/nuPlan/PufferDrive \ From d9120671656dc81ff11e86934f37a020ed056396 Mon Sep 17 00:00:00 2001 From: Bernhard Jaeger Date: Fri, 3 Jul 2026 15:19:14 +0200 Subject: [PATCH 09/35] Update truth values. --- bernhard_configs.txt | 12 +++++------- scripts/0_train_model_slurm.sh | 12 ++++++------ 2 files changed, 11 insertions(+), 13 deletions(-) diff --git a/bernhard_configs.txt b/bernhard_configs.txt index 525fa1f0cc..239ea1d9b0 100644 --- a/bernhard_configs.txt +++ b/bernhard_configs.txt @@ -47,18 +47,18 @@ torchrun --standalone --nnodes=1 --nproc-per-node=1 --max_restarts=0 --start-met --policy.ego-input-size 128 \ --policy.lane-input-size 128 \ --policy.partner-input-size 128 \ - --policy.shared-network false \ + --policy.shared-network False \ --policy.traffic-control-input-size 128 \ --train.checkpoint-interval 50 \ - --train.compile true \ - --train.normalize-rewards false \ + --train.compile True \ + --train.normalize-rewards False \ --train.precision bfloat16 \ --train.seed 4 \ --train.update-epochs 3 \ --policy.boundary-input-size 128 \ --tb \ - --eval.validation-replay.enabled false \ - --eval.behaviors-full-dir.enabled false + --eval.validation-replay.enabled False \ + --eval.behaviors-full-dir.enabled False uv run convert --py123d_path /home/jaeger/ordnung/data/nuPlan/123D --output /home/jaeger/ordnung/data/nuPlan/PufferDrive --preset nuplan @@ -138,8 +138,6 @@ env.traffic-control-scope=0 \ policyboundary-input-size=128 \ render-obs=0 \ -pufferl.py: error: unrecognized arguments: --env.traffic-control-scope 0 --policyboundary-input-size 128 --render-obs 0 true - Usage=puffer [--load-model-path LOAD_MODEL_PATH] [--load-id LOAD_ID] [--render-mode {auto,human,ansi,rgb_array,raylib,None}] [--video-path VIDEO_PATH] [--num_scenarios NUM_SCENARIOS] [--render RENDER] [--agent_index [AGENT_INDEX ...]] [--save-frames SAVE_FRAMES] [--gif-path GIF_PATH] [--fps FPS] [--max-runs MAX_RUNS] [--wandb] [--wandb-project WANDB_PROJECT] [--wandb-group WANDB_GROUP] [--run-name RUN_NAME] [--neptune] diff --git a/scripts/0_train_model_slurm.sh b/scripts/0_train_model_slurm.sh index abf8ea34b5..86393a5c74 100644 --- a/scripts/0_train_model_slurm.sh +++ b/scripts/0_train_model_slurm.sh @@ -6,15 +6,15 @@ #SBATCH --gres gpu:8 #SBATCH --mem 1024G #SBATCH --cpus-per-task 224 -#SBATCH --output /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/k_006_longrun_300B/log_%a_%A.out -#SBATCH --error /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/k_006_longrun_300B/log_%a_%A.err +#SBATCH --output /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/k_007_longrun_300B/log_%a_%A.out +#SBATCH --error /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/k_007_longrun_300B/log_%a_%A.err #SBATCH --partition kyutai # print info about current job echo "START TIME: $(date)" start=$(date +%s) -export RUN_NAME=k_006_longrun_300B +export RUN_NAME=k_007_longrun_300B echo ${RUN_NAME} # TODO could try to tune these. 1 Is probably best since Puffer parallelizes across all cores. @@ -63,11 +63,11 @@ torchrun --standalone --nnodes=1 --nproc-per-node=8 --max_restarts=0 --start-met --policy.ego-input-size 128 \ --policy.lane-input-size 128 \ --policy.partner-input-size 128 \ - --policy.shared-network false \ + --policy.shared-network False \ --policy.traffic-control-input-size 128 \ --train.checkpoint-interval 50 \ - --train.compile true \ - --train.normalize-rewards false \ + --train.compile True \ + --train.normalize-rewards False \ --train.precision bfloat16 \ --train.seed 4 \ --train.update-epochs 3 \ From 37fd7e2ee987fdba61ead3aeaffe6f7d77dbba39 Mon Sep 17 00:00:00 2001 From: Bernhard Jaeger Date: Fri, 3 Jul 2026 15:32:21 +0200 Subject: [PATCH 10/35] 008 try with 48 envs --- bernhard_configs.txt | 4 ++-- scripts/0_train_model_slurm.sh | 9 ++++----- 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/bernhard_configs.txt b/bernhard_configs.txt index 239ea1d9b0..9ea5c3ef53 100644 --- a/bernhard_configs.txt +++ b/bernhard_configs.txt @@ -12,8 +12,8 @@ puffer train puffer_drive --train.learning_rate 0.001 torchrun --standalone --nnodes=1 --nproc-per-node=1 --max_restarts=0 --start-method spawn -m pufferlib.pufferl train puffer_drive --wandb --train.data-dir /home/jaeger/ordnung/internal/PufferDrive/experiments/debug_000 \ --wandb --wandb-project pufferdrive --wandb-group cluster --train.checkpoint-interval 1000 --run-name debug_000 \ --train.total-timesteps 30000000000 \ - --train.max-minibatch-size 65536 \ - --train.minibatch-size 65536 \ + --train.max-minibatch-size 32768 \ + --train.minibatch-size 32768 \ --vec.num-envs 16 \ --vec.num-workers 16 \ --eval.validation-replay.env.map-dir /lustre/scwpod02/client/kyutai/kesai/data/nuPlan/PufferDrive \ diff --git a/scripts/0_train_model_slurm.sh b/scripts/0_train_model_slurm.sh index 86393a5c74..315c4bfd60 100644 --- a/scripts/0_train_model_slurm.sh +++ b/scripts/0_train_model_slurm.sh @@ -6,15 +6,15 @@ #SBATCH --gres gpu:8 #SBATCH --mem 1024G #SBATCH --cpus-per-task 224 -#SBATCH --output /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/k_007_longrun_300B/log_%a_%A.out -#SBATCH --error /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/k_007_longrun_300B/log_%a_%A.err +#SBATCH --output /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/k_008_longrun_300B/log_%a_%A.out +#SBATCH --error /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/k_008_longrun_300B/log_%a_%A.err #SBATCH --partition kyutai # print info about current job echo "START TIME: $(date)" start=$(date +%s) -export RUN_NAME=k_007_longrun_300B +export RUN_NAME=k_008_longrun_300B echo ${RUN_NAME} # TODO could try to tune these. 1 Is probably best since Puffer parallelizes across all cores. @@ -30,8 +30,7 @@ torchrun --standalone --nnodes=1 --nproc-per-node=8 --max_restarts=0 --start-met --train.total-timesteps 300000000000 \ --train.max-minibatch-size 128000 \ --train.minibatch-size 128000 \ - --vec.num-envs 24 \ - --vec.num-workers 24 \ + --vec.num-envs 48 \ --eval.validation-replay.env.map-dir /lustre/scwpod02/client/kyutai/kesai/data/nuPlan/PufferDrive \ --eval.behaviors-full-dir.env.map-dir /lustre/scwpod02/client/kyutai/kesai/data/nuPlan/PufferDrive \ --env.dt 0.3 \ From 503ea4094be4bdb5ae18257feac56c288e5549b0 Mon Sep 17 00:00:00 2001 From: Bernhard Jaeger Date: Fri, 3 Jul 2026 15:45:33 +0200 Subject: [PATCH 11/35] Try gradient accumulation 2 --- bernhard_configs.txt | 2 +- scripts/0_train_model_slurm.sh | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/bernhard_configs.txt b/bernhard_configs.txt index 9ea5c3ef53..682b7c5d1e 100644 --- a/bernhard_configs.txt +++ b/bernhard_configs.txt @@ -13,7 +13,7 @@ torchrun --standalone --nnodes=1 --nproc-per-node=1 --max_restarts=0 --start-met --wandb --wandb-project pufferdrive --wandb-group cluster --train.checkpoint-interval 1000 --run-name debug_000 \ --train.total-timesteps 30000000000 \ --train.max-minibatch-size 32768 \ - --train.minibatch-size 32768 \ + --train.minibatch-size 131072 \ --vec.num-envs 16 \ --vec.num-workers 16 \ --eval.validation-replay.env.map-dir /lustre/scwpod02/client/kyutai/kesai/data/nuPlan/PufferDrive \ diff --git a/scripts/0_train_model_slurm.sh b/scripts/0_train_model_slurm.sh index 315c4bfd60..5eefe3bf3d 100644 --- a/scripts/0_train_model_slurm.sh +++ b/scripts/0_train_model_slurm.sh @@ -28,7 +28,7 @@ source .venv/bin/activate torchrun --standalone --nnodes=1 --nproc-per-node=8 --max_restarts=0 --start-method spawn -m pufferlib.pufferl train puffer_drive --wandb --train.data-dir /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/${RUN_NAME} \ --wandb --wandb-project pufferdrive --wandb-group cluster --train.checkpoint-interval 1000 --run-name ${RUN_NAME} \ --train.total-timesteps 300000000000 \ - --train.max-minibatch-size 128000 \ + --train.max-minibatch-size 64000 \ --train.minibatch-size 128000 \ --vec.num-envs 48 \ --eval.validation-replay.env.map-dir /lustre/scwpod02/client/kyutai/kesai/data/nuPlan/PufferDrive \ From 230ffdb249ce3b4741a3ce0950b9577d830295ee Mon Sep 17 00:00:00 2001 From: Bernhard Jaeger Date: Fri, 3 Jul 2026 16:15:52 +0200 Subject: [PATCH 12/35] half again --- scripts/0_train_model_slurm.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/0_train_model_slurm.sh b/scripts/0_train_model_slurm.sh index 5eefe3bf3d..3f841bb528 100644 --- a/scripts/0_train_model_slurm.sh +++ b/scripts/0_train_model_slurm.sh @@ -28,7 +28,7 @@ source .venv/bin/activate torchrun --standalone --nnodes=1 --nproc-per-node=8 --max_restarts=0 --start-method spawn -m pufferlib.pufferl train puffer_drive --wandb --train.data-dir /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/${RUN_NAME} \ --wandb --wandb-project pufferdrive --wandb-group cluster --train.checkpoint-interval 1000 --run-name ${RUN_NAME} \ --train.total-timesteps 300000000000 \ - --train.max-minibatch-size 64000 \ + --train.max-minibatch-size 32000 \ --train.minibatch-size 128000 \ --vec.num-envs 48 \ --eval.validation-replay.env.map-dir /lustre/scwpod02/client/kyutai/kesai/data/nuPlan/PufferDrive \ From cc93eeacfb26945f12dfc0e757c179a8f452a12c Mon Sep 17 00:00:00 2001 From: bjaeger Date: Mon, 20 Jul 2026 17:22:04 +0000 Subject: [PATCH 13/35] update script --- scripts/0_build_puffer.sh | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/scripts/0_build_puffer.sh b/scripts/0_build_puffer.sh index 92020d5e25..dd1bbaeae0 100644 --- a/scripts/0_build_puffer.sh +++ b/scripts/0_build_puffer.sh @@ -6,9 +6,9 @@ #SBATCH --gres=gpu:1 #SBATCH --mem=64G #SBATCH --cpus-per-task=24 -#SBATCH --output=/lustre/scwpod02/client/kyutai/kesai/bernhard/carla_closed_loop/results/logs/eval_server_%a_%A.out -#SBATCH --error=/lustre/scwpod02/client/kyutai/kesai/bernhard/carla_closed_loop/results/logs/eval_server_%a_%A.err -#SBATCH --partition=kyutai +#SBATCH --output=/home/bjaeger/logs/eval_server_%a_%A.out +#SBATCH --error=/home/bjaeger/logs/eval_server_%a_%A.err +#SBATCH --partition=dev # print info about current job echo "START TIME: $(date)" From c9464697ce26ead37871697d64d9196780bf2c9f Mon Sep 17 00:00:00 2001 From: bernhard_jaeger Date: Tue, 21 Jul 2026 14:20:18 +0000 Subject: [PATCH 14/35] Remove val --- pufferlib/config/ocean/drive.ini | 40 ++++++++++++------------- scripts/1_train_model_slurm.sh | 51 ++++++++++++++++++++++++++++++++ 2 files changed, 71 insertions(+), 20 deletions(-) create mode 100644 scripts/1_train_model_slurm.sh diff --git a/pufferlib/config/ocean/drive.ini b/pufferlib/config/ocean/drive.ini index 1da2eaf509..8501127fc8 100644 --- a/pufferlib/config/ocean/drive.ini +++ b/pufferlib/config/ocean/drive.ini @@ -299,21 +299,21 @@ eval.num_scenarios = 250 eval.export_episode_csv = true eval.verify_coverage = true -[eval.validation_replay] -inherits = "validation_defaults" -type = "multi_scenario" -enabled = true -render = true -render_backend = "triage_html" -env.simulation_mode = "replay" -env.control_mode = "control_sdc_only" -env.map_dir = "/lustre/scwpod02/client/kyutai/kesai/data/nuPlan/PufferDrive" -env.num_maps = 250 -env.max_agents_per_env = 64 -env.scenario_length = 200 -env.resample_frequency = 200 -eval.render_num_scenarios = 5 -eval.render_max_steps = 200 +; [eval.validation_replay] +; inherits = "validation_defaults" +; type = "multi_scenario" +; enabled = true +; render = true +; render_backend = "triage_html" +; env.simulation_mode = "replay" +; env.control_mode = "control_sdc_only" +; env.map_dir = "/lustre/scwpod02/client/kyutai/kesai/data/nuPlan/PufferDrive" +; env.num_maps = 250 +; env.max_agents_per_env = 64 +; env.scenario_length = 200 +; env.resample_frequency = 200 +; eval.render_num_scenarios = 5 +; eval.render_max_steps = 200 [eval.validation_gigaflow] inherits = "validation_defaults" @@ -372,11 +372,11 @@ eval.num_scenarios = 50 eval.render_num_scenarios = 2 eval.render_max_steps = 200 -[eval.behaviors_full_dir] -inherits = "behaviors_defaults" -type = "behavior_class" -enabled = true -env.map_dir = "/lustre/scwpod02/client/kyutai/kesai/data/nuPlan/PufferDrive" +; [eval.behaviors_full_dir] +; inherits = "behaviors_defaults" +; type = "behavior_class" +; enabled = true +; env.map_dir = "/lustre/scwpod02/client/kyutai/kesai/data/nuPlan/PufferDrive" ; [eval.behaviors_hard_stop] ; inherits = "behaviors_defaults" diff --git a/scripts/1_train_model_slurm.sh b/scripts/1_train_model_slurm.sh new file mode 100644 index 0000000000..f00da483c6 --- /dev/null +++ b/scripts/1_train_model_slurm.sh @@ -0,0 +1,51 @@ +#!/bin/bash +#SBATCH --job-name train_puffer +#SBATCH --ntasks 1 +#SBATCH --nodes 1 +#SBATCH --time 2-00:00 +#SBATCH --gres gpu:8 +#SBATCH --mem=1007G +#SBATCH --cpus-per-task 144 +#SBATCH --output /home/bjaeger/PufferDrive/experiments/k_nightly_0001/log_%a_%A.out +#SBATCH --error /home/bjaeger/PufferDrive/experiments/k_nightly_0001/log_%a_%A.err +#SBATCH --partition dev + +# print info about current job +echo "START TIME: $(date)" +start=$(date +%s) + +export RUN_NAME=k_nightly_0001 +echo ${RUN_NAME} + +# TODO could try to tune these. 1 Is probably best since Puffer parallelizes across all cores. +export NUMEXPR_NUM_THREADS=1 +export MKL_NUM_THREADS=1 +export VECLIB_MAXIMUM_THREADS=1 +export OPENBLAS_NUM_THREADS=1 +export OMP_NUM_THREADS=1 + +source .venv/bin/activate +torchrun --standalone --nnodes=1 --nproc-per-node=8 --max_restarts=0 --start-method spawn \ + -m pufferlib.pufferl train puffer_drive \ + wandb=True \ + wandb_project=nightly-multi-long \ + wandb_group=emerge_ \ + train.data_dir=/home/bjaeger/PufferDrive/experiments/${RUN_NAME} \ + env.map_dir=/home/bjaeger/PufferDrive/pufferlib/resources/drive/binaries/carla \ + train.name=${RUN_NAME} \ + train.total_timesteps=100000000000 \ + vec.num_envs=16 \ + +eval.map_dir=/home/bjaeger/data/nuPlan/PufferDrive \ + train.compile=True \ + train.max_minibatch_size=131072 \ + train.minibatch_size=131072 \ + train.precision=bfloat16 \ + env.num_agents=4096 \ + tb=True + + + +end=$(date +%s) +runtime=$((end-start)) +echo "END TIME: $(date)" +echo "Runtime: ${runtime}" From 4ad1e7c07d4a2fe727f3194c0c8f6e8479488788 Mon Sep 17 00:00:00 2001 From: Bernhard Jaeger Date: Tue, 21 Jul 2026 21:12:42 +0200 Subject: [PATCH 15/35] Added a parameter such that the model can be discrete if the environment is continuous. --- notebooks/06_architecture.py | 3 + notebooks/notebook_utils.py | 1 + pufferlib/config/ocean/drive.ini | 32 +++++----- pufferlib/ocean/benchmark/evaluator.py | 24 ++++++-- pufferlib/ocean/benchmark/evaluators/base.py | 56 ++++++++++++----- pufferlib/ocean/torch.py | 41 ++++++++++++- pufferlib/pufferl.py | 44 ++++++++++---- pufferlib/pytorch.py | 19 ++++-- scripts/0_build_puffer.sh | 6 +- scripts/0_train_model_slurm.sh | 10 ++-- scripts/2_train_model_multinode_slurm.sh | 63 ++++++++++++++++++++ 11 files changed, 237 insertions(+), 62 deletions(-) create mode 100644 scripts/2_train_model_multinode_slurm.sh diff --git a/notebooks/06_architecture.py b/notebooks/06_architecture.py index 879dc42f84..687f89cb69 100644 --- a/notebooks/06_architecture.py +++ b/notebooks/06_architecture.py @@ -40,6 +40,7 @@ BACKBONE_ACTIVATION = "gelu" BACKBONE_LAYER_NORM = False MASK_PADDED_FEATURES = False +ACTION_TYPE = "discrete" env, obs, info = make_drive_env() @@ -64,6 +65,7 @@ backbone_layer_norm=BACKBONE_LAYER_NORM, shared_network=SHARED_NETWORK, mask_padded_features=MASK_PADDED_FEATURES, + action_type=ACTION_TYPE, ).to(device) print(f"Device: {device}") @@ -519,6 +521,7 @@ def count_params(module): "backbone_layer_norm": False, "shared_network": True, "mask_padded_features": False, + "action_type": "discrete" } results = [] diff --git a/notebooks/notebook_utils.py b/notebooks/notebook_utils.py index 6db9e01d44..530cc8554d 100644 --- a/notebooks/notebook_utils.py +++ b/notebooks/notebook_utils.py @@ -112,6 +112,7 @@ "backbone_layer_norm": False, "shared_network": True, "mask_padded_features": False, + "action_type": "discrete" } diff --git a/pufferlib/config/ocean/drive.ini b/pufferlib/config/ocean/drive.ini index 7934f8266e..acb25484fa 100644 --- a/pufferlib/config/ocean/drive.ini +++ b/pufferlib/config/ocean/drive.ini @@ -36,6 +36,8 @@ critic_hidden_size = 512 critic_num_layers = 0 ; Dual or shared actor-critic backbone shared_network = True +; Actions type that the policy uses - options: "discrete", "continuous", will be converted to continuous for the simulator. +action_type = "discrete" [rnn] input_size = 512 @@ -290,21 +292,21 @@ eval.num_scenarios = 250 eval.export_episode_csv = true eval.verify_coverage = true -[eval.validation_replay] -inherits = "validation_defaults" -type = "multi_scenario" -enabled = true -render = true -render_backend = "triage_html" -env.simulation_mode = "replay" -env.control_mode = "control_sdc_only" -env.map_dir = "/lustre/scwpod02/client/kyutai/kesai/data/nuPlan/PufferDrive" -env.num_maps = 250 -env.max_agents_per_env = 64 -env.scenario_length = 200 -env.resample_frequency = 200 -eval.render_num_scenarios = 5 -eval.render_max_steps = 200 +; [eval.validation_replay] +; inherits = "validation_defaults" +; type = "multi_scenario" +; enabled = true +; render = true +; render_backend = "triage_html" +; env.simulation_mode = "replay" +; env.control_mode = "control_sdc_only" +; env.map_dir = "/lustre/scwpod02/client/kyutai/kesai/data/nuPlan/PufferDrive" +; env.num_maps = 250 +; env.max_agents_per_env = 64 +; env.scenario_length = 200 +; env.resample_frequency = 200 +; eval.render_num_scenarios = 5 +; eval.render_max_steps = 200 [eval.validation_gigaflow] inherits = "validation_defaults" diff --git a/pufferlib/ocean/benchmark/evaluator.py b/pufferlib/ocean/benchmark/evaluator.py index 404476e4e7..f970eff41e 100644 --- a/pufferlib/ocean/benchmark/evaluator.py +++ b/pufferlib/ocean/benchmark/evaluator.py @@ -88,6 +88,8 @@ def collect_simulated_trajectories(self, args, puffer_env, policy): "id": np.zeros((num_agents, self.num_rollouts, self.sim_steps), dtype=np.int32), } + env_continuous = isinstance(puffer_env.single_action_space, pufferlib.spaces.Box) + for rollout_idx in range(self.num_rollouts): print(f"\rCollecting rollout {rollout_idx + 1}/{self.num_rollouts}...", end="", flush=True) obs, info = puffer_env.reset() @@ -98,6 +100,7 @@ def collect_simulated_trajectories(self, args, puffer_env, policy): lstm_c=torch.zeros(num_agents, policy.hidden_size, device=device), ) + # TODO find out how to test this code path for time_idx in range(self.sim_steps): # Get global state agent_state = driver.get_global_agent_state() @@ -111,13 +114,19 @@ def collect_simulated_trajectories(self, args, puffer_env, policy): with torch.no_grad(): ob_tensor = torch.as_tensor(obs).to(device) logits, value = policy.forward_eval(ob_tensor, state) - action, logprob, _ = pufferlib.pytorch.sample_logits(logits) + action, logprob, _, cont_action = pufferlib.pytorch.sample_logits(logits, action=None, deterministic=False, env_continuous=env_continuous, policy=policy) action_np = action.cpu().numpy().reshape(puffer_env.action_space.shape) if isinstance(logits, torch.distributions.Normal): action_np = np.clip(action_np, puffer_env.action_space.low, puffer_env.action_space.high) - obs, _, _, _, _ = puffer_env.step(action_np) + + if env_continuous and not policy.is_continuous: # TODO check with trajectory + cont_action = cont_action.cpu().numpy().reshape(puffer_env.action_space.shape) + obs, _, _, _, _ = puffer_env.step(cont_action) + else: + obs, _, _, _, _ = puffer_env.step(action_np) + return trajectories @@ -659,6 +668,7 @@ def rollout(self, args, puffer_env, policy): total_steps = (scenario_length - init_steps + 1) * num_maps obs, info = puffer_env.reset() + env_continuous = isinstance(puffer_env.single_action_space, pufferlib.spaces.Box) state = {} if args["train"]["use_rnn"]: state = dict( @@ -671,13 +681,19 @@ def rollout(self, args, puffer_env, policy): with torch.no_grad(): ob_tensor = torch.as_tensor(obs).to(device) logits, value = policy.forward_eval(ob_tensor, state) - action, logprob, _ = pufferlib.pytorch.sample_logits(logits) + action, logprob, _, cont_action = pufferlib.pytorch.sample_logits(logits, action=None, deterministic=False, env_continuous=env_continuous, policy=policy) action_np = action.cpu().numpy().reshape(puffer_env.action_space.shape) if isinstance(logits, torch.distributions.Normal): action_np = np.clip(action_np, puffer_env.action_space.low, puffer_env.action_space.high) - obs, rewards, dones, truncs, info_list = puffer_env.step(action_np) + + if env_continuous and not policy.is_continuous: # TODO check with trajectory + cont_action = cont_action.cpu().numpy().reshape(puffer_env.action_space.shape) + obs, rewards, dones, truncs, info_list = puffer_env.step(cont_action) + else: + obs, rewards, dones, truncs, info_list = puffer_env.step(action_np) + if info_list: all_infos.extend(info_list) # Stop once we've collected at least one info per bin to avoid diff --git a/pufferlib/ocean/benchmark/evaluators/base.py b/pufferlib/ocean/benchmark/evaluators/base.py index de51ece27f..61cfcf861a 100644 --- a/pufferlib/ocean/benchmark/evaluators/base.py +++ b/pufferlib/ocean/benchmark/evaluators/base.py @@ -150,16 +150,22 @@ def _run_rollout_loop(self, vecenv, policy, args) -> dict: stall_limit = 3 * scenario_length if scenario_length else 0 last_progress = 0 stall_steps = 0 + env_continuous = isinstance(vecenv.single_action_space, pufferlib.spaces.Box) while not self._should_stop(args, infos_collected, steps): with torch.no_grad(): ob_t = torch.as_tensor(obs).to(device) logits, _ = policy.forward_eval(ob_t, state) - action, _, _ = pufferlib.pytorch.sample_logits(logits, deterministic=True) + action, _, _, cont_action = pufferlib.pytorch.sample_logits(logits, action=None, deterministic=True, env_continuous=env_continuous, policy=policy) + + if env_continuous and not policy.is_continuous: # TODO check with trajectory + cont_action = cont_action.cpu().numpy().reshape(vecenv.action_space.shape) + obs, _, terminals, truncations, infos = vecenv.step(cont_action) + else: action = action.cpu().numpy().reshape(vecenv.action_space.shape) - if isinstance(logits, torch.distributions.Normal): - action = np.clip(action, vecenv.action_space.low, vecenv.action_space.high) + if isinstance(logits, torch.distributions.Normal): + action = np.clip(action, vecenv.action_space.low, vecenv.action_space.high) + obs, _, terminals, truncations, infos = vecenv.step(action) - obs, _, terminals, truncations, infos = vecenv.step(action) for d in self._flatten_infos(infos): if isinstance(d, dict) and d.get("summary_type") == "completed_episode": self._episode_rows.append(d) @@ -512,17 +518,21 @@ def _render_pass_html(self, vecenv, policy, args) -> list: if state: state["lstm_h"].zero_() state["lstm_c"].zero_() - + env_continuous = isinstance(vecenv.single_action_space, pufferlib.spaces.Box) for _ in range(max_steps * num_scenarios): with torch.no_grad(): ob_t = torch.as_tensor(ob).to(device) logits, _ = policy.forward_eval(ob_t, state) - action, _, _ = pufferlib.pytorch.sample_logits(logits, deterministic=True) - action = action.cpu().numpy().reshape(vec.action_space.shape) - if isinstance(logits, torch.distributions.Normal): - action = np.clip(action, vec.action_space.low, vec.action_space.high) + action, _, _, cont_action = pufferlib.pytorch.sample_logits(logits, actions=None, deterministic=True, env_continuous=env_continuous, policy=policy) - ob, _, terminals, truncations, infos = vec.step(action) + if env_continuous and not policy.is_continuous: # TODO check with trajectory + cont_action = cont_action.cpu().numpy().reshape(vecenv.action_space.shape) + ob, _, terminals, truncations, infos = vec.step(cont_action) + else: + action = action.cpu().numpy().reshape(vec.action_space.shape) + if isinstance(logits, torch.distributions.Normal): + action = np.clip(action, vec.action_space.low, vec.action_space.high) + ob, _, terminals, truncations, infos = vec.step(action) if state: done = np.asarray(terminals).astype(bool) | np.asarray(truncations).astype(bool) @@ -620,6 +630,7 @@ def _render_pass_obs(self, vecenv, policy, args) -> list: ) try: state = self._init_lstm_state(vec.observation_space.shape[0], policy, device, args) + env_continuous = isinstance(vec.single_action_space, pufferlib.spaces.Box) while scenarios_done < num_scenarios: ob, _ = vec.reset() scenarios = vec.get_state() @@ -677,7 +688,7 @@ def _render_pass_obs(self, vecenv, policy, args) -> list: ob_t = torch.as_tensor(ob).to(device) logits, value = policy.forward_eval(ob_t, state) pool_outputs = pool_method(ob_t, state) if pool_method is not None else {} - action, logprob, entropy = pufferlib.pytorch.sample_logits(logits, deterministic=True) + action, logprob, entropy, cont_action = pufferlib.pytorch.sample_logits(logits, actions=None, deterministic=True, env_continuous=env_continuous, policy=policy) raw_action = action.cpu().numpy().reshape(vec.action_space.shape) pool_outputs = {k: v.cpu().numpy().astype(np.int16, copy=False) for k, v in pool_outputs.items()} if pool_hist is None and pool_outputs: @@ -748,7 +759,13 @@ def _render_pass_obs(self, vecenv, policy, args) -> list: np.asarray(policy_outputs[start_obs_index:end_obs_index], dtype=np.float32).copy() ) start_obs_index = end_obs_index - ob, _, _, _, step_infos = vec.step(clipped_action) + + if env_continuous and not policy.is_continuous: # TODO check with trajectory + cont_action = cont_action.cpu().numpy().reshape(vec.action_space.shape) + ob, _, _, _, step_infos = vec.step(cont_action) + else: + ob, _, _, _, step_infos = vec.step(clipped_action) + for d in self._flatten_infos(step_infos): if isinstance(d, dict) and d.get("summary_type") == "completed_episode": batch_summary = d @@ -838,6 +855,7 @@ def _render_view(self, vecenv, target_env, policy, args, view_idx, out_dir, view # The trailing `.mp4` is exact, so e.g. `*_epoch7_step12_bev.mp4` # matches bev files but not the bare sim_state ones. view_glob = f"*{view_suffix}.mp4" + env_continuous = isinstance(vecenv.single_action_space, pufferlib.spaces.Box) try: state = self._init_lstm_state(num_agents, policy, device, args) scenarios_processed = 0 @@ -856,11 +874,17 @@ def _render_view(self, vecenv, target_env, policy, args, view_idx, out_dir, view with torch.no_grad(): ob_t = torch.as_tensor(ob).to(device) logits, _ = policy.forward_eval(ob_t, state) - action, _, _ = pufferlib.pytorch.sample_logits(logits, deterministic=True) + action, _, _, cont_action = pufferlib.pytorch.sample_logits(logits, actions=None, deterministic=True, env_continuous=env_continuous, policy=policy) + + if env_continuous and not policy.is_continuous: # TODO check with trajectory + cont_action = cont_action.cpu().numpy().reshape(vecenv.action_space.shape) + ob, _, _, _, _ = vecenv.step(cont_action) + else: action = action.cpu().numpy().reshape(vecenv.action_space.shape) - if isinstance(logits, torch.distributions.Normal): - action = np.clip(action, vecenv.action_space.low, vecenv.action_space.high) - ob, _, _, _, _ = vecenv.step(action) + if isinstance(logits, torch.distributions.Normal): + action = np.clip(action, vecenv.action_space.low, vecenv.action_space.high) + ob, _, _, _, _ = vecenv.step(action) + for e in range(to_render): target_env.render(env_idx=e, view_mode=view_idx) for e in range(to_render): diff --git a/pufferlib/ocean/torch.py b/pufferlib/ocean/torch.py index 6bcf1533c5..e7ac70ffcb 100644 --- a/pufferlib/ocean/torch.py +++ b/pufferlib/ocean/torch.py @@ -325,8 +325,35 @@ def __init__( backbone_layer_norm: bool, shared_network: bool, mask_padded_features: bool, + action_type: str, ): super().__init__() + + self.register_buffer("JERK_LONG", torch.tensor((-15.0, -4.0, 0.0, 4.0), dtype=torch.float32, requires_grad=False), persistent=False) + self.register_buffer("JERK_LAT", torch.tensor((-4.0, 0.0, 4.0), dtype=torch.float32, requires_grad=False), persistent=False) + + self.register_buffer("ACCELERATION_VALUES", torch.tensor((-4.0000, -2.6670, -1.3330, -0.0000, 1.3330, 2.6670, 4.0000), dtype=torch.float32, requires_grad=False), persistent=False) + self.register_buffer("STEERING_VALUES", torch.tensor((-0.667, -0.500, -0.333, -0.167, 0.000, 0.167, 0.333, 0.500, 0.667), dtype=torch.float32, requires_grad=False), persistent=False) + + if env.dynamics_model == "jerk": + action_long, action_lat = self.JERK_LONG, self.JERK_LAT + elif env.dynamics_model == "classic": + action_long, action_lat = self.ACCELERATION_VALUES, self.STEERING_VALUES + else: + raise ValueError(f"Unsupported dynamics model: {env.dynamics_model}") + + # Precompute the [-1, 1] continuous action per discrete choice by inverting the + # sim's continuous scaling (drive.h). Constant → done once. Longitudinal is + # asymmetric (braking / |t[0]|, accel / t[-1]); lateral symmetric. The symmetric + # classic table collapses both branches to / t[-1]. + long_norm = torch.where( + action_long < 0.0, + action_long / -action_long[0], + action_long / action_long[-1], + ) + lat_norm = action_lat / action_lat[-1] + self.register_buffer("action_long_norm", long_norm, persistent=False) + self.register_buffer("action_lat_norm", lat_norm, persistent=False) # Configuration flags from policy kwargs self.shared_network = shared_network @@ -362,11 +389,11 @@ def __init__( self.critic_backbone = DriveBackbone(**backbone_args) # Setup action and value heads - self.is_continuous = isinstance(env.single_action_space, pufferlib.spaces.Box) + self.is_continuous = action_type == "continuous" # TODO Check if what the `"trajectory"`, `"trajectory_frenet"`, `"trajectory_jerk" features do and if they are considered continuous. if self.is_continuous: self.atn_dim = (env.single_action_space.shape[0],) * 2 else: - self.atn_dim = env.single_action_space.nvec.tolist() + self.atn_dim = [self.action_long_norm.numel() * self.action_lat_norm.numel()] # n-layer MLP for actor head (num_layers = number of hidden layers) backbone_out_dim = self.actor_backbone.out_dim @@ -449,3 +476,13 @@ def decode_actions(self, hidden): value = self.critic_head(hidden) return action, value + + def discrete_actions_to_continuous(self, actions): + actions = actions.long() + num_lat = self.action_lat_norm.numel() + long_val = self.action_long_norm[actions // num_lat] + lat_val = self.action_lat_norm[actions % num_lat] + + return torch.stack([long_val, lat_val], dim=-1) + + diff --git a/pufferlib/pufferl.py b/pufferlib/pufferl.py index eeb8364053..4faf85db65 100644 --- a/pufferlib/pufferl.py +++ b/pufferlib/pufferl.py @@ -128,8 +128,15 @@ def __init__(self, config, vecenv, policy, logger=None): # Vecenv info vecenv.async_reset(seed) + + self.env_continuous = isinstance(vecenv.single_action_space, pufferlib.spaces.Box) obs_space = vecenv.single_observation_space - atn_space = vecenv.single_action_space + if self.env_continuous and not policy.is_continuous: + action_shape = (len(policy.atn_dim),) + action_dtype = torch.int32 + else: + action_shape = vecenv.single_action_space + action_dtype = pufferlib.pytorch.numpy_to_torch_dtype_dict[vecenv.single_action_space.dtype] total_agents = vecenv.num_agents self.total_agents = total_agents @@ -171,9 +178,9 @@ def __init__(self, config, vecenv, policy, logger=None): self.actions = torch.zeros( segments, horizon, - *atn_space.shape, + *action_shape, device=device, - dtype=pufferlib.pytorch.numpy_to_torch_dtype_dict[atn_space.dtype], + dtype=action_dtype, ) self.values = torch.zeros(segments, horizon, device=device) self.logprobs = torch.zeros(segments, horizon, device=device) @@ -399,7 +406,8 @@ def evaluate(self): logits, value = self.policy.forward_eval(o_device, state) logits = logits_to_float(logits) - action, logprob, _ = pufferlib.pytorch.sample_logits(logits) + action, logprob, _, cont_action = pufferlib.pytorch.sample_logits(logits, action=None, deterministic=False, + env_continuous=self.env_continuous, policy=self.policy) if config["normalize_rewards"]: r = torch.sign(r) * torch.log1p(torch.abs(r)) @@ -441,10 +449,7 @@ def evaluate(self): self.ep_lengths[env_id] = 0 self.free_idx += num_full self.full_rows += num_full - - action = action.cpu().numpy() - if isinstance(logits, torch.distributions.Normal): - action = np.clip(action, self.vecenv.action_space.low, self.vecenv.action_space.high) + profile("eval_misc", epoch) for i in info: @@ -457,7 +462,15 @@ def evaluate(self): self.stats[k].append(v) profile("env", epoch) - self.vecenv.send(action) + + if self.env_continuous and not self.policy.is_continuous: # TODO check with trajectory + cont_action = cont_action.cpu().numpy() + self.vecenv.send(cont_action.squeeze(0)) + else: + action = action.cpu().numpy() + if isinstance(logits, torch.distributions.Normal): + action = np.clip(action, self.vecenv.action_space.low, self.vecenv.action_space.high) + self.vecenv.send(action) profile("eval_misc", epoch) self.free_idx = self.total_agents @@ -560,7 +573,7 @@ def _ppo_loss(self, mb_obs, mb_actions, mb_logprobs, mb_values, mb_returns, mb_a logits, newvalue = self.policy(mb_obs, state) logits = logits_to_float(logits) newvalue = newvalue.float() - _, newlogprob, entropy = pufferlib.pytorch.sample_logits(logits, action=mb_actions) + _, newlogprob, entropy, _ = pufferlib.pytorch.sample_logits(logits, action=mb_actions) newlogprob = newlogprob.float().view_as(mb_logprobs) newvalue = newvalue.view_as(mb_returns) @@ -572,6 +585,7 @@ def _ppo_loss(self, mb_obs, mb_actions, mb_logprobs, mb_values, mb_returns, mb_a approx_kl = ((ratio - 1) - logratio).mean() clipfrac = ((ratio - 1.0).abs() > config["clip_coef"]).float().mean() + # TODO fix multi-gpu bug mb_adv = (mb_adv - mb_adv.mean()) / (mb_adv.std(unbiased=False) + 1e-8) if adv_weights is not None: mb_adv = adv_weights * mb_adv @@ -1830,17 +1844,23 @@ def mine_failures(env_name, args=None): obs_arr, *_ = vecenv.recv() pbar_total = num_episodes pbar_done = 0 + env_continuous = isinstance(vecenv.single_action_space, pufferlib.spaces.Box) print(f"[mine_failures] target episodes={num_episodes} output={output_dir} score_threshold={score_threshold}") while pbar_done < num_episodes: with torch.no_grad(): o_t = torch.as_tensor(obs_arr).to(device) state = {"reward": None, "done": None, "env_id": None, "mask": None} logits, _ = policy.forward_eval(o_t, state) - action, _, _ = pufferlib.pytorch.sample_logits(logits) + action, _, _, cont_action = pufferlib.pytorch.sample_logits(logits, action=None, deterministic=False, env_continuous=env_continuous, policy=policy) action = action.cpu().numpy() if action.ndim == 1 and len(vecenv.single_action_space.shape) >= 1: action = action.reshape(-1, *vecenv.single_action_space.shape) - vecenv.send(action) + + if env_continuous and not policy.is_continuous: # TODO check with trajectory + cont_action = cont_action.cpu().numpy() + vecenv.send(cont_action.squeeze(0)) + else: + vecenv.send(action) obs_arr, _, _, _, infos, *_ = vecenv.recv() for info in infos: if not isinstance(info, dict): diff --git a/pufferlib/pytorch.py b/pufferlib/pytorch.py index 2ec10a5c5f..466b1126ec 100644 --- a/pufferlib/pytorch.py +++ b/pufferlib/pytorch.py @@ -187,7 +187,7 @@ def entropy_probs(logits, probs): return -p_log_p.sum(-1) -def sample_logits(logits, action=None, deterministic=False): +def sample_logits(logits, action=None, deterministic=False, env_continuous=None, policy=None): # TODO discrete continuous is_discrete = isinstance(logits, torch.Tensor) if isinstance(logits, torch.distributions.Normal): batch = logits.loc.shape[0] @@ -198,7 +198,7 @@ def sample_logits(logits, action=None, deterministic=False): log_probs = logits.log_prob(action.view(batch, -1)).sum(1) logits_entropy = logits.entropy().view(batch, -1).sum(1) - return action, log_probs, logits_entropy + return action, log_probs, logits_entropy, None elif is_discrete: logits = logits.unsqueeze(0) # TODO: Double check this @@ -211,9 +211,10 @@ def sample_logits(logits, action=None, deterministic=False): normalized_logits = logits - logits.logsumexp(dim=-1, keepdim=True) probs = logits_to_probs(logits) + # TODO this jumps in here and makes the action determinisitc if action is None: if deterministic: - action = torch.argmax(probs, -1) + action = torch.argmax(probs, -1) # TODO weighted average else: probs = torch.nan_to_num(probs, 1e-8, 1e-8, 1e-8) action = torch.multinomial(probs.reshape(-1, probs.shape[-1]), 1, replacement=True).int() @@ -226,7 +227,15 @@ def sample_logits(logits, action=None, deterministic=False): logprob = log_prob(normalized_logits, action) logits_entropy = entropy(normalized_logits).sum(0) + if env_continuous and not policy.is_continuous: + cont_actions = policy.discrete_actions_to_continuous(action) + + if is_discrete: + return action.squeeze(0), logprob.squeeze(0), logits_entropy.squeeze(0), cont_actions.squeeze(0) + + return action.T, logprob.sum(0), logits_entropy, cont_actions + if is_discrete: - return action.squeeze(0), logprob.squeeze(0), logits_entropy.squeeze(0) + return action.squeeze(0), logprob.squeeze(0), logits_entropy.squeeze(0), None - return action.T, logprob.sum(0), logits_entropy + return action.T, logprob.sum(0), logits_entropy, None diff --git a/scripts/0_build_puffer.sh b/scripts/0_build_puffer.sh index 92020d5e25..dd1bbaeae0 100644 --- a/scripts/0_build_puffer.sh +++ b/scripts/0_build_puffer.sh @@ -6,9 +6,9 @@ #SBATCH --gres=gpu:1 #SBATCH --mem=64G #SBATCH --cpus-per-task=24 -#SBATCH --output=/lustre/scwpod02/client/kyutai/kesai/bernhard/carla_closed_loop/results/logs/eval_server_%a_%A.out -#SBATCH --error=/lustre/scwpod02/client/kyutai/kesai/bernhard/carla_closed_loop/results/logs/eval_server_%a_%A.err -#SBATCH --partition=kyutai +#SBATCH --output=/home/bjaeger/logs/eval_server_%a_%A.out +#SBATCH --error=/home/bjaeger/logs/eval_server_%a_%A.err +#SBATCH --partition=dev # print info about current job echo "START TIME: $(date)" diff --git a/scripts/0_train_model_slurm.sh b/scripts/0_train_model_slurm.sh index 3f841bb528..c59dfa2688 100644 --- a/scripts/0_train_model_slurm.sh +++ b/scripts/0_train_model_slurm.sh @@ -4,11 +4,11 @@ #SBATCH --nodes 1 #SBATCH --time 2-00:00 #SBATCH --gres gpu:8 -#SBATCH --mem 1024G -#SBATCH --cpus-per-task 224 -#SBATCH --output /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/k_008_longrun_300B/log_%a_%A.out -#SBATCH --error /lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/k_008_longrun_300B/log_%a_%A.err -#SBATCH --partition kyutai +#SBATCH --mem=1007G +#SBATCH --cpus-per-task 144 +#SBATCH --output /home/bjaeger/PufferDrive/experiments/k_008_longrun_300B/log_%a_%A.out +#SBATCH --error /home/bjaeger/PufferDrive/experiments/k_008_longrun_300B/log_%a_%A.err +#SBATCH --partition dev # print info about current job echo "START TIME: $(date)" diff --git a/scripts/2_train_model_multinode_slurm.sh b/scripts/2_train_model_multinode_slurm.sh new file mode 100644 index 0000000000..420c0ac36c --- /dev/null +++ b/scripts/2_train_model_multinode_slurm.sh @@ -0,0 +1,63 @@ +#!/bin/bash +#SBATCH --job-name train_puffer +#SBATCH --nodes 2 # Number of nodes requested +#SBATCH --ntasks-per-node 1 # Run 1 srun task per node (which fires up torchrun) +#SBATCH --gres gpu:8 # GPUs per node +#SBATCH --cpus-per-task 144 +#SBATCH --mem=1007G +#SBATCH --time 2-00:00 +#SBATCH --output /home/bjaeger/PufferDrive/experiments/k_nightly_0002/log_%a_%A.out +#SBATCH --error /home/bjaeger/PufferDrive/experiments/k_nightly_0002/log_%a_%A.err +#SBATCH --partition dev + +# Set up PyTorch Distributed Rendezvous parameters from Slurm variables +export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1) +export MASTER_PORT=29500 +export WORLD_SIZE=$SLURM_JOB_NUM_NODES + +echo "START TIME: $(date)" +echo "Master node: ${MASTER_ADDR}:${MASTER_PORT}" +echo "Total nodes: ${WORLD_SIZE}" +start=$(date +%s) + +export RUN_NAME=k_nightly_0002 +echo ${RUN_NAME} + +# Thread limit limits CPU thrashing across worker environments +export NUMEXPR_NUM_THREADS=1 +export MKL_NUM_THREADS=1 +export VECLIB_MAXIMUM_THREADS=1 +export OPENBLAS_NUM_THREADS=1 +export OMP_NUM_THREADS=1 + +source .venv/bin/activate + +# Execute torchrun across all nodes using srun +srun torchrun \ + --nnodes=${SLURM_JOB_NUM_NODES} \ + --nproc-per-node=8 \ + --rdzv_id=${SLURM_JOB_ID} \ + --rdzv_backend=c10d \ + --rdzv_endpoint="${MASTER_ADDR}:${MASTER_PORT}" \ + --max_restarts=0 \ + --start-method spawn \ + -m pufferlib.pufferl train puffer_drive \ + wandb=True \ + wandb_project=nightly-multi-long \ + wandb_group=emerge_ \ + train.data_dir=/home/bjaeger/PufferDrive/experiments/${RUN_NAME} \ + env.map_dir=/home/bjaeger/PufferDrive/pufferlib/resources/drive/binaries/carla \ + train.name=${RUN_NAME} \ + train.total_timesteps=100000000000 \ + vec.num_envs=16 \ + +eval.map_dir=/home/bjaeger/data/nuPlan/PufferDrive \ + train.compile=True \ + train.max_minibatch_size=131072 \ + train.minibatch_size=131072 \ + train.precision=bfloat16 \ + tb=True + +end=$(date +%s) +runtime=$((end-start)) +echo "END TIME: $(date)" +echo "Runtime: ${runtime}" From b209fbd52b69b7b1f41a6790b5571637895c2dfe Mon Sep 17 00:00:00 2001 From: Bernhard Jaeger Date: Tue, 21 Jul 2026 21:13:43 +0200 Subject: [PATCH 16/35] start changes --- notebooks/05_inference.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/notebooks/05_inference.py b/notebooks/05_inference.py index 0df79f5295..a38288a0f5 100644 --- a/notebooks/05_inference.py +++ b/notebooks/05_inference.py @@ -76,8 +76,8 @@ logits_list, value = policy(obs_tensor) # Sample actions -action, logprob, ent = sample_logits(logits_list) -action_det, _, _ = sample_logits(logits_list, deterministic=True) +action, logprob, ent, _ = sample_logits(logits_list) +action_det, _, _, _ = sample_logits(logits_list, deterministic=True) # TODO print(f"Value: mean={value.mean():.4f}, std={value.std():.4f}, range=[{value.min():.4f}, {value.max():.4f}]") print(f"Entropy: mean={ent.mean():.4f}, std={ent.std():.4f}") @@ -143,7 +143,7 @@ def run_rollout(env, policy, deterministic=False, horizon=HORIZON): obs_t = torch.FloatTensor(obs).to(device) with torch.no_grad(): logits_list, val = policy(obs_t) - act, logp, entr = sample_logits(logits_list, deterministic=deterministic) + act, logp, entr, _ = sample_logits(logits_list, deterministic=deterministic) buffers["obs"][t] = obs buffers["actions"][t] = act.cpu().numpy().reshape(N) if act.dim() > 1 else act.cpu().numpy() From bbed9f13378c495f2ae63984d916cd07f123e661 Mon Sep 17 00:00:00 2001 From: Bernhard Jaeger Date: Wed, 22 Jul 2026 16:01:17 +0200 Subject: [PATCH 17/35] fixed 05_inference --- notebooks/05_inference.py | 24 +- pufferlib/config/puffer_drive.yaml | 1 + scripts/baseline_10G/config-reward-dense.yaml | 1 + .../baseline_10G/config-reward-sparse.yaml | 1 + scripts/cluster_configs/nightly_best.yaml | 1 + .../single_agent_speed_run.yaml | 1 + scripts/config-reward-dense.yaml | 403 ++++++++++++++++++ scripts/config-reward-sparse.yaml | 403 ++++++++++++++++++ weights/mimolette/config.yaml | 1 + 9 files changed, 828 insertions(+), 8 deletions(-) create mode 100644 scripts/config-reward-dense.yaml create mode 100644 scripts/config-reward-sparse.yaml diff --git a/notebooks/05_inference.py b/notebooks/05_inference.py index a0b9c06591..7175123799 100644 --- a/notebooks/05_inference.py +++ b/notebooks/05_inference.py @@ -25,6 +25,7 @@ from pufferlib.ocean.drive import binding from pufferlib.ocean.torch import Drive as DrivePolicy from pufferlib.pytorch import sample_logits +import pufferlib.spaces from notebooks.notebook_utils import COEF_NAMES, EGO_LABELS, MAP_DIR, load_notebook_config, zero_actions CHECKPOINT_PATH = "" @@ -54,8 +55,9 @@ policy.load_state_dict(sd) print(f"Loaded checkpoint: {CHECKPOINT_PATH}") -is_continuous = policy.is_continuous -ACT_SHAPE = (N, len(env.single_action_space.nvec)) if not is_continuous else (N, env.single_action_space.shape[0]) +is_continuous = policy.is_continuous # policy output space +env_continuous = isinstance(env.single_action_space, pufferlib.spaces.Box) # env action space +ACT_SHAPE = (N, env.single_action_space.shape[0]) if env_continuous else (N, len(env.single_action_space.nvec)) print(f"Policy on {device}, params: {sum(p.numel() for p in policy.parameters()):,}") print(f"Obs shape: {obs.shape}, Action space: {env.single_action_space}") @@ -76,8 +78,8 @@ logits_list, value = policy(obs_tensor) # Sample actions -action, logprob, ent, _ = sample_logits(logits_list) -action_det, _, _, _ = sample_logits(logits_list, deterministic=True) # TODO +action, logprob, ent, cont_action = sample_logits(logits_list, env_continuous=env_continuous, policy=policy) +action_det, _, _, _ = sample_logits(logits_list, deterministic=True, env_continuous=env_continuous, policy=policy) print(f"Value: mean={value.mean():.4f}, std={value.std():.4f}, range=[{value.min():.4f}, {value.max():.4f}]") print(f"Entropy: mean={ent.mean():.4f}, std={ent.std():.4f}") @@ -143,7 +145,9 @@ def run_rollout(env, policy, deterministic=False, horizon=HORIZON): obs_t = torch.FloatTensor(obs).to(device) with torch.no_grad(): logits_list, val = policy(obs_t) - act, logp, entr, _ = sample_logits(logits_list, deterministic=deterministic) + act, logp, entr, cont_act = sample_logits( + logits_list, deterministic=deterministic, env_continuous=env_continuous, policy=policy + ) buffers["obs"][t] = obs buffers["actions"][t] = act.cpu().numpy().reshape(N) if act.dim() > 1 else act.cpu().numpy() @@ -156,8 +160,12 @@ def run_rollout(env, policy, deterministic=False, horizon=HORIZON): buffers["positions_x"][t] = gstate["x"] buffers["positions_y"][t] = gstate["y"] - # Step env - env_actions = act.cpu().numpy().reshape(ACT_SHAPE) + # Step env. A discrete policy on a continuous env steps with the decoded + # continuous action; otherwise the (discrete or continuous) action itself. + if env_continuous and not is_continuous: + env_actions = cont_act.cpu().numpy().reshape(N, -1) + else: + env_actions = act.cpu().numpy().reshape(ACT_SHAPE) obs, rew, term, trunc, info = env.step(env_actions) buffers["rewards"][t] = rew buffers["terminals"][t] = term @@ -1049,7 +1057,7 @@ def unpack_all_timesteps(bufs, agent_idx): # %% # Compute action probs over time for tracked agent (stochastic rollout) -n_actions = env.single_action_space.nvec[0] if not is_continuous else 1 +n_actions = policy.atn_dim[0] if not is_continuous else 1 action_probs_time = np.zeros((HORIZON, n_actions)) for t in range(HORIZON): obs_t = torch.FloatTensor(buf_stoch["obs"][t : t + 1, TRACKED_AGENT : TRACKED_AGENT + 1][0]).to(device) diff --git a/pufferlib/config/puffer_drive.yaml b/pufferlib/config/puffer_drive.yaml index 54da056a0f..22d0eef9c4 100644 --- a/pufferlib/config/puffer_drive.yaml +++ b/pufferlib/config/puffer_drive.yaml @@ -195,6 +195,7 @@ policy: critic_num_layers: 0 # Dual or shared actor-critic backbone shared_network: true + action_type: discrete rnn: input_size: 512 hidden_size: 512 diff --git a/scripts/baseline_10G/config-reward-dense.yaml b/scripts/baseline_10G/config-reward-dense.yaml index f9eed82808..77922f036d 100644 --- a/scripts/baseline_10G/config-reward-dense.yaml +++ b/scripts/baseline_10G/config-reward-dense.yaml @@ -324,6 +324,7 @@ policy: partner_input_size: 128 shared_network: true traffic_control_input_size: 128 + action_type: discrete policy_name: Drive render: 0 render_mode: auto diff --git a/scripts/baseline_10G/config-reward-sparse.yaml b/scripts/baseline_10G/config-reward-sparse.yaml index 645d7c9523..77497998a0 100644 --- a/scripts/baseline_10G/config-reward-sparse.yaml +++ b/scripts/baseline_10G/config-reward-sparse.yaml @@ -324,6 +324,7 @@ policy: partner_input_size: 128 shared_network: true traffic_control_input_size: 128 + action_type: discrete policy_name: Drive render: 0 render_mode: auto diff --git a/scripts/cluster_configs/nightly_best.yaml b/scripts/cluster_configs/nightly_best.yaml index 3b603aa660..b5fce31016 100644 --- a/scripts/cluster_configs/nightly_best.yaml +++ b/scripts/cluster_configs/nightly_best.yaml @@ -95,6 +95,7 @@ policy.backbone_num_layers: 3 policy.critic_hidden_size: 256 policy.critic_num_layers: 1 policy.shared_network: false +policy.action_type: discrete # Training — large minibatch, compiled bfloat16 train.total_timesteps: 10_000_000_000 diff --git a/scripts/cluster_configs/single_agent_speed_run.yaml b/scripts/cluster_configs/single_agent_speed_run.yaml index befabd3a1b..b8531b93ec 100644 --- a/scripts/cluster_configs/single_agent_speed_run.yaml +++ b/scripts/cluster_configs/single_agent_speed_run.yaml @@ -95,6 +95,7 @@ policy.backbone_num_layers: 3 policy.critic_hidden_size: 256 policy.critic_num_layers: 1 policy.shared_network: false +policy.action_type: discrete # Traffic lights fully off: not observed, not scored, no reward penalty. env.traffic_light_behavior: 0 diff --git a/scripts/config-reward-dense.yaml b/scripts/config-reward-dense.yaml new file mode 100644 index 0000000000..77922f036d --- /dev/null +++ b/scripts/config-reward-dense.yaml @@ -0,0 +1,403 @@ +agent_index: null +controlled_exp: + train: + ent_coef: + values: + - 0.01 + - 0.005 + learning_rate: + values: + - 0.001 + - 0.003 + - 0.01 +env: + action_type: discrete + collision_behavior: 1 + compute_eval_metrics: false + control_mode: control_vehicles + dt: 0.1 + dynamics_model: jerk + goal_on_lane: true + goal_radius: 2.0 + goal_speed: 1000.0 + inactive_agent_threshold: 0.4 + init_mode: create_all_valid + init_step: 0 + map_dir: "pufferlib/resources/drive/binaries/carla" + max_agents_per_env: 80 + max_waypoint_spacing: 40.0 + min_agents_per_env: 1 + min_waypoint_spacing: 20.0 + non_sdc_controller: policy + non_vehicle_controller: auto + num_agents: 2048 + num_maps: 8 + num_target_waypoints: 3 + obs_dropout_boundary: 0.0 + obs_dropout_lane: 0.0 + obs_norm_goal_offset_m: 120.0 + obs_norm_road_seg_length_m: 10.0 + obs_norm_road_seg_width_m: 5.0 + obs_norm_veh_length_m: 10.0 + obs_norm_veh_width_m: 5.0 + obs_norm_xy_offset_m: 120.0 + obs_range_partner_m: 120.0 + obs_range_road_behind_m: 30.0 + obs_range_road_front_m: 120.0 + obs_range_road_side_m: 40.0 + obs_range_traffic_control_m: 120.0 + obs_slots_boundary_n: 30 + obs_slots_lane_n: 50 + obs_slots_partners_n: 12 + obs_slots_traffic_controls_n: 4 + offroad_behavior: 1 + partner_blindness_prob: 0.0 + partner_blindness_trigger_prob: 0.0 + phantom_braking_duration: 10 + phantom_braking_prob: 0.0 + phantom_braking_trigger_prob: 0.0 + resample_frequency: 102400 + reward_ade: 0.0 + reward_center_bias: 0.0 + reward_collision: 1.5 + reward_comfort: 0.5 + reward_conditioning: false + reward_goal: 0.2 + reward_lane_align: 0.5 + reward_lane_center: 0.05 + reward_offroad: 1.5 + reward_overspeed: 0.1 + reward_randomization: false + reward_reverse: 0.005 + reward_stop_line: 1.0 + reward_timestep: 0.0 + reward_vel_align: 1.0 + reward_velocity: 0.1 + scenario_length: 1024 + sdc_controller: policy + simulation_mode: gigaflow + spawn_initial_speed: 0.0 + goal_regen_mode: finite + goal_source: route + termination_mode: 1 + traffic_light_behavior: 1 + use_map_cache: 1 +env_name: puffer_drive +eval: + behaviors_defaults: + clean: 'true' + enabled: 'false' + env: + control_mode: control_sdc_only + init_mode: create_all_valid + obs_slots_partners_n: 32 + scenario_length: 201 + simulation_mode: replay + eval: + num_scenarios: 50 + render_max_steps: 200 + render_num_scenarios: 2 + interval: 250 + mode: inline + render: 'true' + render_views: + - sim_state + - bev + behaviors_full_dir: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/nuplan_mini_train_bins + inherits: behaviors_defaults + type: behavior_class + behaviors_hard_stop: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/hard_stop + inherits: behaviors_defaults + type: behavior_class + behaviors_highway_straight: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/highway_straight + inherits: behaviors_defaults + type: behavior_class + behaviors_lane_change: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/lane_change + inherits: behaviors_defaults + type: behavior_class + behaviors_merge: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/merge + inherits: behaviors_defaults + type: behavior_class + behaviors_parked_cars: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/parked_cars + inherits: behaviors_defaults + type: behavior_class + behaviors_roundabout: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/roundabout + inherits: behaviors_defaults + type: behavior_class + behaviors_stopped_traffic: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/stopped_traffic + inherits: behaviors_defaults + type: behavior_class + behaviors_traffic_light_green: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/traffic_light_green + inherits: behaviors_defaults + type: behavior_class + behaviors_traffic_light_stop: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/traffic_light_stop + inherits: behaviors_defaults + type: behavior_class + behaviors_unprotected_left: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/unprotected_left + inherits: behaviors_defaults + type: behavior_class + behaviors_unprotected_right: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/unprotected_right + inherits: behaviors_defaults + type: behavior_class + dnf_triage: + enabled: 'false' + env: + map_dir: pufferlib/resources/drive/binaries/carla/opendrive__Town10HD.bin + max_agents_per_env: 1 + min_agents_per_env: 1 + num_maps: 1 + resample_frequency: 500 + scenario_length: 500 + simulation_mode: gigaflow + eval: + num_scenarios: 32 + render_max_steps: 300 + render_num_scenarios: 16 + inherits: validation_defaults + render: 'true' + render_backend: triage_html + type: multi_scenario + validation_defaults: + clean: 'true' + enabled: 'true' + env: + collision_behavior: 1 + eval_mode: 1 + goal_speed: 3.0 + num_agents: 1024 + obs_dropout_boundary: 0.0 + obs_dropout_lane: 0.0 + obs_slots_boundary_n: 80 + obs_slots_lane_n: 80 + offroad_behavior: 1 + reward_ade: 0.0 + reward_collision: 3.0 + reward_comfort: 0.05 + reward_goal: 1.0 + reward_lane_align: 0.025 + reward_lane_center: 0.0038 + reward_offroad: 3.0 + reward_overspeed: 0.05 + reward_randomization: false + reward_reverse: 0.005 + reward_stop_line: 1.0 + reward_timestep: 2.5e-05 + reward_velocity: 0.0025 + goal_source: route + goal_regen_mode: finite + termination_mode: 0 + traffic_light_behavior: 0 + eval: + export_episode_csv: 'true' + num_scenarios: 250 + verify_coverage: 'true' + interval: 250 + mode: inline + validation_gigaflow: + enabled: 'true' + env: + map_dir: pufferlib/resources/drive/binaries/carla + max_agents_per_env: 40 + min_agents_per_env: 40 + num_agents: 1024 + num_maps: 8 + resample_frequency: 500 + scenario_length: 500 + simulation_mode: gigaflow + eval: + render_max_steps: 300 + render_num_scenarios: 8 + inherits: validation_defaults + render: 'true' + render_backend: egl + render_views: + - sim_state + - bev + type: multi_scenario + validation_replay: + enabled: 'true' + env: + control_mode: control_sdc_only + map_dir: /scratch/ev2237/data/nuplan/nuplan_mini_train_bins + max_agents_per_env: 64 + num_maps: 250 + resample_frequency: 200 + scenario_length: 200 + simulation_mode: replay + eval: + render_max_steps: 200 + render_num_scenarios: 5 + inherits: validation_defaults + render: 'true' + render_backend: triage_html + type: multi_scenario + wosac: + clean: 'true' + enabled: 'false' + env: + control_mode: control_wosac + goal_radius: 2.0 + init_mode: create_all_valid + init_step: 10 + eval: + wosac_aggregate_results: 'true' + wosac_num_agents: 256 + wosac_num_rollouts: 32 + wosac_sanity_check: 'false' + interval: 500 + mode: subprocess + render: 'false' + type: wosac +eval_simulation: null +fps: 15 +gif_path: eval.gif +git: + commit_hash: 9d3fcc5c09db50b9674a374b8c3404b6349c4cb5 +load_id: null +load_model_path: null +local_rank: 0 +max_runs: 200 +max_suggestion_cost: 3600 +mine: + num_episodes: 100 + output_dir: '' + render: 'true' + score_threshold: -inf +neptune: false +neptune_name: pufferai +neptune_project: ablations +no_model_upload: {} +num_scenarios: 3 +package: ocean +policy: + actor_hidden_size: 512 + actor_num_layers: 0 + backbone_activation: gelu + backbone_hidden_size: 512 + backbone_layer_norm: false + backbone_num_layers: 2 + boundary_input_size: 128 + context_input_size: 128 + critic_hidden_size: 512 + critic_num_layers: 0 + ego_input_size: 128 + encoder_activation: relu + encoder_layer_norm: true + lane_input_size: 128 + mask_padded_features: false + partner_input_size: 128 + shared_network: true + traffic_control_input_size: 128 + action_type: discrete +policy_name: Drive +render: 0 +render_mode: auto +rnn: + hidden_size: 512 + input_size: 512 +rnn_name: null +run_name: rew-dense-seed4 +save_frames: 0 +sweep: + downsample: 10 + goal: maximize + method: Protein + metric: score +tag: null +tb: false +train: + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1.0e-08 + adv_filter_ewma_beta: 0.25 + adv_filter_threshold_scale: 0.01 + adv_sampling_prio_alpha: 0.8499999999999999 + adv_sampling_prio_beta0: 0.8499999999999999 + amp: true + anneal_lr: true + batch_size: auto + bptt_horizon: 128 + checkpoint_interval: 50 + clip_coef: 0.2 + compile: true + compile_fullgraph: false + compile_mode: default + cpu_offload: false + data_dir: /pufferdrive/training_output/ + device: cuda + ent_coef: 0.01 + gae_lambda: 0.95 + gamma: 0.999 + learning_rate: 0.0005 + max_grad_norm: 0.5 + max_minibatch_size: 98304 + minibatch_size: 98304 + name: pufferai + normalize_rewards: true + obs_only: true + optimizer: adamw + precision: bfloat16 + project: ablations + render: false + render_interval: 1000 + render_map: none + resume_state_path: null + seed: 4 + show_grid: false + show_human_logs: true + show_lasers: false + torch_deterministic: false + total_timesteps: 10000000000 + update_epochs: 2 + use_rnn: false + vf_clip_coef: null + vf_coef: 0.5 + vtrace_c_clip: 1 + vtrace_rho_clip: 1 +vec: + backend: Multiprocessing + batch_size: auto + num_envs: 24 + num_workers: auto + seed: 42 + zero_copy: true +video_path: videos +wandb: true +wandb_group: pr-smoke-runs +wandb_project: pr-smoke-runs diff --git a/scripts/config-reward-sparse.yaml b/scripts/config-reward-sparse.yaml new file mode 100644 index 0000000000..77497998a0 --- /dev/null +++ b/scripts/config-reward-sparse.yaml @@ -0,0 +1,403 @@ +agent_index: null +controlled_exp: + train: + ent_coef: + values: + - 0.01 + - 0.005 + learning_rate: + values: + - 0.001 + - 0.003 + - 0.01 +env: + action_type: discrete + collision_behavior: 1 + compute_eval_metrics: false + control_mode: control_vehicles + dt: 0.1 + dynamics_model: jerk + goal_on_lane: true + goal_radius: 2.0 + goal_speed: 1000.0 + inactive_agent_threshold: 0.4 + init_mode: create_all_valid + init_step: 0 + map_dir: "pufferlib/resources/drive/binaries/carla" + max_agents_per_env: 80 + max_waypoint_spacing: 40.0 + min_agents_per_env: 1 + min_waypoint_spacing: 20.0 + non_sdc_controller: policy + non_vehicle_controller: auto + num_agents: 2048 + num_maps: 8 + num_target_waypoints: 3 + obs_dropout_boundary: 0.0 + obs_dropout_lane: 0.0 + obs_norm_goal_offset_m: 120.0 + obs_norm_road_seg_length_m: 10.0 + obs_norm_road_seg_width_m: 5.0 + obs_norm_veh_length_m: 10.0 + obs_norm_veh_width_m: 5.0 + obs_norm_xy_offset_m: 120.0 + obs_range_partner_m: 120.0 + obs_range_road_behind_m: 30.0 + obs_range_road_front_m: 120.0 + obs_range_road_side_m: 40.0 + obs_range_traffic_control_m: 120.0 + obs_slots_boundary_n: 30 + obs_slots_lane_n: 50 + obs_slots_partners_n: 12 + obs_slots_traffic_controls_n: 4 + offroad_behavior: 1 + partner_blindness_prob: 0.0 + partner_blindness_trigger_prob: 0.0 + phantom_braking_duration: 10 + phantom_braking_prob: 0.0 + phantom_braking_trigger_prob: 0.0 + resample_frequency: 102400 + reward_ade: 0.0 + reward_center_bias: 0.0 + reward_collision: 1.5 + reward_comfort: 0.05 + reward_conditioning: false + reward_goal: 0.5 + reward_lane_align: 0.025 + reward_lane_center: 0.0038 + reward_offroad: 1.5 + reward_overspeed: 0.05 + reward_randomization: false + reward_reverse: 0.005 + reward_stop_line: 1.0 + reward_timestep: 0.0 + reward_vel_align: 1.0 + reward_velocity: 0.0025 + scenario_length: 1024 + sdc_controller: policy + simulation_mode: gigaflow + spawn_initial_speed: 0.0 + goal_source: route + goal_regen_mode: finite + termination_mode: 1 + traffic_light_behavior: 1 + use_map_cache: 1 +env_name: puffer_drive +eval: + behaviors_defaults: + clean: 'true' + enabled: 'false' + env: + control_mode: control_sdc_only + init_mode: create_all_valid + obs_slots_partners_n: 32 + scenario_length: 201 + simulation_mode: replay + eval: + num_scenarios: 50 + render_max_steps: 200 + render_num_scenarios: 2 + interval: 250 + mode: inline + render: 'true' + render_views: + - sim_state + - bev + behaviors_full_dir: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/nuplan_mini_train_bins + inherits: behaviors_defaults + type: behavior_class + behaviors_hard_stop: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/hard_stop + inherits: behaviors_defaults + type: behavior_class + behaviors_highway_straight: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/highway_straight + inherits: behaviors_defaults + type: behavior_class + behaviors_lane_change: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/lane_change + inherits: behaviors_defaults + type: behavior_class + behaviors_merge: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/merge + inherits: behaviors_defaults + type: behavior_class + behaviors_parked_cars: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/parked_cars + inherits: behaviors_defaults + type: behavior_class + behaviors_roundabout: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/roundabout + inherits: behaviors_defaults + type: behavior_class + behaviors_stopped_traffic: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/stopped_traffic + inherits: behaviors_defaults + type: behavior_class + behaviors_traffic_light_green: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/traffic_light_green + inherits: behaviors_defaults + type: behavior_class + behaviors_traffic_light_stop: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/traffic_light_stop + inherits: behaviors_defaults + type: behavior_class + behaviors_unprotected_left: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/unprotected_left + inherits: behaviors_defaults + type: behavior_class + behaviors_unprotected_right: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/unprotected_right + inherits: behaviors_defaults + type: behavior_class + dnf_triage: + enabled: 'false' + env: + map_dir: pufferlib/resources/drive/binaries/carla/opendrive__Town10HD.bin + max_agents_per_env: 1 + min_agents_per_env: 1 + num_maps: 1 + resample_frequency: 500 + scenario_length: 500 + simulation_mode: gigaflow + eval: + num_scenarios: 32 + render_max_steps: 300 + render_num_scenarios: 16 + inherits: validation_defaults + render: 'true' + render_backend: triage_html + type: multi_scenario + validation_defaults: + clean: 'true' + enabled: 'true' + env: + collision_behavior: 1 + eval_mode: 1 + goal_speed: 3.0 + num_agents: 1024 + obs_dropout_boundary: 0.0 + obs_dropout_lane: 0.0 + obs_slots_boundary_n: 80 + obs_slots_lane_n: 80 + offroad_behavior: 1 + reward_ade: 0.0 + reward_collision: 3.0 + reward_comfort: 0.05 + reward_goal: 1.0 + reward_lane_align: 0.025 + reward_lane_center: 0.0038 + reward_offroad: 3.0 + reward_overspeed: 0.05 + reward_randomization: false + reward_reverse: 0.005 + reward_stop_line: 1.0 + reward_timestep: 2.5e-05 + reward_velocity: 0.0025 + goal_source: route + goal_regen_mode: finite + termination_mode: 0 + traffic_light_behavior: 0 + eval: + export_episode_csv: 'true' + num_scenarios: 250 + verify_coverage: 'true' + interval: 250 + mode: inline + validation_gigaflow: + enabled: 'true' + env: + map_dir: pufferlib/resources/drive/binaries/carla + max_agents_per_env: 40 + min_agents_per_env: 40 + num_agents: 1024 + num_maps: 8 + resample_frequency: 500 + scenario_length: 500 + simulation_mode: gigaflow + eval: + render_max_steps: 300 + render_num_scenarios: 8 + inherits: validation_defaults + render: 'true' + render_backend: egl + render_views: + - sim_state + - bev + type: multi_scenario + validation_replay: + enabled: 'true' + env: + control_mode: control_sdc_only + map_dir: /scratch/ev2237/data/nuplan/nuplan_mini_train_bins + max_agents_per_env: 64 + num_maps: 250 + resample_frequency: 200 + scenario_length: 200 + simulation_mode: replay + eval: + render_max_steps: 200 + render_num_scenarios: 5 + inherits: validation_defaults + render: 'true' + render_backend: triage_html + type: multi_scenario + wosac: + clean: 'true' + enabled: 'false' + env: + control_mode: control_wosac + goal_radius: 2.0 + init_mode: create_all_valid + init_step: 10 + eval: + wosac_aggregate_results: 'true' + wosac_num_agents: 256 + wosac_num_rollouts: 32 + wosac_sanity_check: 'false' + interval: 500 + mode: subprocess + render: 'false' + type: wosac +eval_simulation: null +fps: 15 +gif_path: eval.gif +git: + commit_hash: 9d3fcc5c09db50b9674a374b8c3404b6349c4cb5 +load_id: null +load_model_path: null +local_rank: 0 +max_runs: 200 +max_suggestion_cost: 3600 +mine: + num_episodes: 100 + output_dir: '' + render: 'true' + score_threshold: -inf +neptune: false +neptune_name: pufferai +neptune_project: ablations +no_model_upload: {} +num_scenarios: 3 +package: ocean +policy: + actor_hidden_size: 512 + actor_num_layers: 0 + backbone_activation: gelu + backbone_hidden_size: 512 + backbone_layer_norm: false + backbone_num_layers: 2 + boundary_input_size: 128 + context_input_size: 128 + critic_hidden_size: 512 + critic_num_layers: 0 + ego_input_size: 128 + encoder_activation: relu + encoder_layer_norm: true + lane_input_size: 128 + mask_padded_features: false + partner_input_size: 128 + shared_network: true + traffic_control_input_size: 128 + action_type: discrete +policy_name: Drive +render: 0 +render_mode: auto +rnn: + hidden_size: 512 + input_size: 512 +rnn_name: null +run_name: rew-sparse-seed4 +save_frames: 0 +sweep: + downsample: 10 + goal: maximize + method: Protein + metric: score +tag: null +tb: false +train: + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1.0e-08 + adv_filter_ewma_beta: 0.25 + adv_filter_threshold_scale: 0.01 + adv_sampling_prio_alpha: 0.8499999999999999 + adv_sampling_prio_beta0: 0.8499999999999999 + amp: true + anneal_lr: true + batch_size: auto + bptt_horizon: 128 + checkpoint_interval: 50 + clip_coef: 0.2 + compile: true + compile_fullgraph: false + compile_mode: default + cpu_offload: false + data_dir: /pufferdrive/training_output/ + device: cuda + ent_coef: 0.01 + gae_lambda: 0.95 + gamma: 0.999 + learning_rate: 0.0005 + max_grad_norm: 0.5 + max_minibatch_size: 98304 + minibatch_size: 98304 + name: pufferai + normalize_rewards: true + obs_only: true + optimizer: adamw + precision: bfloat16 + project: ablations + render: false + render_interval: 1000 + render_map: none + resume_state_path: null + seed: 4 + show_grid: false + show_human_logs: true + show_lasers: false + torch_deterministic: false + total_timesteps: 10000000000 + update_epochs: 2 + use_rnn: false + vf_clip_coef: null + vf_coef: 0.5 + vtrace_c_clip: 1 + vtrace_rho_clip: 1 +vec: + backend: Multiprocessing + batch_size: auto + num_envs: 24 + num_workers: auto + seed: 42 + zero_copy: true +video_path: videos +wandb: true +wandb_group: pr-smoke-runs +wandb_project: pr-smoke-runs diff --git a/weights/mimolette/config.yaml b/weights/mimolette/config.yaml index 3d63068e63..bf8291429c 100644 --- a/weights/mimolette/config.yaml +++ b/weights/mimolette/config.yaml @@ -141,6 +141,7 @@ policy: partner_input_size: 256 shared_network: false traffic_control_input_size: 128 + action_type: discrete policy_name: Drive render: 0 render_mode: matplotlib From aa56963ff90bb6c211c0f544bcf33f5efeafec7d Mon Sep 17 00:00:00 2001 From: Bernhard Jaeger Date: Wed, 22 Jul 2026 18:37:08 +0200 Subject: [PATCH 18/35] Added the option to change between mean mode and sample, but kept everything mode for now. --- notebooks/05_inference.py | 12 +++++----- pufferlib/ocean/benchmark/evaluator.py | 4 ++-- pufferlib/ocean/benchmark/evaluators/base.py | 8 +++---- pufferlib/ocean/torch.py | 19 ++++++++++------ pufferlib/pufferl.py | 5 ++--- pufferlib/pytorch.py | 23 ++++++++++++-------- 6 files changed, 40 insertions(+), 31 deletions(-) diff --git a/notebooks/05_inference.py b/notebooks/05_inference.py index 7175123799..8e5538514b 100644 --- a/notebooks/05_inference.py +++ b/notebooks/05_inference.py @@ -24,7 +24,7 @@ from pufferlib.ocean.drive.drive import Drive from pufferlib.ocean.drive import binding from pufferlib.ocean.torch import Drive as DrivePolicy -from pufferlib.pytorch import sample_logits +from pufferlib.pytorch import sample_logits, ACTION_SELECT_MODE, ACTION_SELECT_SAMPLE import pufferlib.spaces from notebooks.notebook_utils import COEF_NAMES, EGO_LABELS, MAP_DIR, load_notebook_config, zero_actions @@ -79,7 +79,7 @@ # Sample actions action, logprob, ent, cont_action = sample_logits(logits_list, env_continuous=env_continuous, policy=policy) -action_det, _, _, _ = sample_logits(logits_list, deterministic=True, env_continuous=env_continuous, policy=policy) +action_det, _, _, _ = sample_logits(logits_list, action_selection=ACTION_SELECT_MODE, env_continuous=env_continuous, policy=policy) print(f"Value: mean={value.mean():.4f}, std={value.std():.4f}, range=[{value.min():.4f}, {value.max():.4f}]") print(f"Entropy: mean={ent.mean():.4f}, std={ent.std():.4f}") @@ -123,7 +123,7 @@ n_tgt_wp = config["env"].get("num_goals", 3) -def run_rollout(env, policy, deterministic=False, horizon=HORIZON): +def run_rollout(env, policy, action_selection=ACTION_SELECT_SAMPLE, horizon=HORIZON): obs, _ = env.reset(seed=42) N = env.num_agents @@ -146,7 +146,7 @@ def run_rollout(env, policy, deterministic=False, horizon=HORIZON): with torch.no_grad(): logits_list, val = policy(obs_t) act, logp, entr, cont_act = sample_logits( - logits_list, deterministic=deterministic, env_continuous=env_continuous, policy=policy + logits_list, action_selection=action_selection, env_continuous=env_continuous, policy=policy ) buffers["obs"][t] = obs @@ -175,9 +175,9 @@ def run_rollout(env, policy, deterministic=False, horizon=HORIZON): print("Running stochastic rollout...") -buf_stoch = run_rollout(env, policy, deterministic=False) +buf_stoch = run_rollout(env, policy, action_selection=ACTION_SELECT_SAMPLE) print("Running deterministic rollout...") -buf_det = run_rollout(env, policy, deterministic=True) +buf_det = run_rollout(env, policy, action_selection=ACTION_SELECT_MODE) for name, buf in [("Stochastic", buf_stoch), ("Deterministic", buf_det)]: print(f"\n--- {name} ---") diff --git a/pufferlib/ocean/benchmark/evaluator.py b/pufferlib/ocean/benchmark/evaluator.py index f970eff41e..e42927c5b7 100644 --- a/pufferlib/ocean/benchmark/evaluator.py +++ b/pufferlib/ocean/benchmark/evaluator.py @@ -114,7 +114,7 @@ def collect_simulated_trajectories(self, args, puffer_env, policy): with torch.no_grad(): ob_tensor = torch.as_tensor(obs).to(device) logits, value = policy.forward_eval(ob_tensor, state) - action, logprob, _, cont_action = pufferlib.pytorch.sample_logits(logits, action=None, deterministic=False, env_continuous=env_continuous, policy=policy) + action, logprob, _, cont_action = pufferlib.pytorch.sample_logits(logits, env_continuous=env_continuous, policy=policy) action_np = action.cpu().numpy().reshape(puffer_env.action_space.shape) if isinstance(logits, torch.distributions.Normal): @@ -681,7 +681,7 @@ def rollout(self, args, puffer_env, policy): with torch.no_grad(): ob_tensor = torch.as_tensor(obs).to(device) logits, value = policy.forward_eval(ob_tensor, state) - action, logprob, _, cont_action = pufferlib.pytorch.sample_logits(logits, action=None, deterministic=False, env_continuous=env_continuous, policy=policy) + action, logprob, _, cont_action = pufferlib.pytorch.sample_logits(logits, env_continuous=env_continuous, policy=policy) action_np = action.cpu().numpy().reshape(puffer_env.action_space.shape) if isinstance(logits, torch.distributions.Normal): diff --git a/pufferlib/ocean/benchmark/evaluators/base.py b/pufferlib/ocean/benchmark/evaluators/base.py index 61cfcf861a..eaa8cbe266 100644 --- a/pufferlib/ocean/benchmark/evaluators/base.py +++ b/pufferlib/ocean/benchmark/evaluators/base.py @@ -155,7 +155,7 @@ def _run_rollout_loop(self, vecenv, policy, args) -> dict: with torch.no_grad(): ob_t = torch.as_tensor(obs).to(device) logits, _ = policy.forward_eval(ob_t, state) - action, _, _, cont_action = pufferlib.pytorch.sample_logits(logits, action=None, deterministic=True, env_continuous=env_continuous, policy=policy) + action, _, _, cont_action = pufferlib.pytorch.sample_logits(logits, action_selection=pufferlib.pytorch.ACTION_SELECT_MODE, env_continuous=env_continuous, policy=policy) if env_continuous and not policy.is_continuous: # TODO check with trajectory cont_action = cont_action.cpu().numpy().reshape(vecenv.action_space.shape) @@ -523,7 +523,7 @@ def _render_pass_html(self, vecenv, policy, args) -> list: with torch.no_grad(): ob_t = torch.as_tensor(ob).to(device) logits, _ = policy.forward_eval(ob_t, state) - action, _, _, cont_action = pufferlib.pytorch.sample_logits(logits, actions=None, deterministic=True, env_continuous=env_continuous, policy=policy) + action, _, _, cont_action = pufferlib.pytorch.sample_logits(logits, action_selection=pufferlib.pytorch.ACTION_SELECT_MODE, env_continuous=env_continuous, policy=policy) if env_continuous and not policy.is_continuous: # TODO check with trajectory cont_action = cont_action.cpu().numpy().reshape(vecenv.action_space.shape) @@ -688,7 +688,7 @@ def _render_pass_obs(self, vecenv, policy, args) -> list: ob_t = torch.as_tensor(ob).to(device) logits, value = policy.forward_eval(ob_t, state) pool_outputs = pool_method(ob_t, state) if pool_method is not None else {} - action, logprob, entropy, cont_action = pufferlib.pytorch.sample_logits(logits, actions=None, deterministic=True, env_continuous=env_continuous, policy=policy) + action, logprob, entropy, cont_action = pufferlib.pytorch.sample_logits(logits, action_selection=pufferlib.pytorch.ACTION_SELECT_MODE, env_continuous=env_continuous, policy=policy) raw_action = action.cpu().numpy().reshape(vec.action_space.shape) pool_outputs = {k: v.cpu().numpy().astype(np.int16, copy=False) for k, v in pool_outputs.items()} if pool_hist is None and pool_outputs: @@ -874,7 +874,7 @@ def _render_view(self, vecenv, target_env, policy, args, view_idx, out_dir, view with torch.no_grad(): ob_t = torch.as_tensor(ob).to(device) logits, _ = policy.forward_eval(ob_t, state) - action, _, _, cont_action = pufferlib.pytorch.sample_logits(logits, actions=None, deterministic=True, env_continuous=env_continuous, policy=policy) + action, _, _, cont_action = pufferlib.pytorch.sample_logits(logits, action_selection=pufferlib.pytorch.ACTION_SELECT_MODE, env_continuous=env_continuous, policy=policy) if env_continuous and not policy.is_continuous: # TODO check with trajectory cont_action = cont_action.cpu().numpy().reshape(vecenv.action_space.shape) diff --git a/pufferlib/ocean/torch.py b/pufferlib/ocean/torch.py index 00d274885b..02d9ab89bb 100644 --- a/pufferlib/ocean/torch.py +++ b/pufferlib/ocean/torch.py @@ -360,6 +360,14 @@ def __init__( self.register_buffer("action_long_norm", long_norm, persistent=False) self.register_buffer("action_lat_norm", lat_norm, persistent=False) + # Joint continuous-action table: row k = the [-1,1] (long, lat) for discrete class k, + # where k = long_idx * num_lat + lat_idx (matches drive.h decode). + num_lat = lat_norm.numel() + num_classes = long_norm.numel() * num_lat + k = torch.arange(num_classes) + action_table = torch.stack([long_norm[k // num_lat], lat_norm[k % num_lat]], dim=-1) # [num_classes, 2] + self.register_buffer("action_table", action_table, persistent=False) + # Configuration flags from policy kwargs self.shared_network = shared_network self.ego_dim = env.ego_features @@ -483,11 +491,8 @@ def decode_actions(self, hidden): return action, value def discrete_actions_to_continuous(self, actions): - actions = actions.long() - num_lat = self.action_lat_norm.numel() - long_val = self.action_long_norm[actions // num_lat] - lat_val = self.action_lat_norm[actions % num_lat] - - return torch.stack([long_val, lat_val], dim=-1) - + return self.action_table[actions.long()] + def discrete_probs_to_continuous_mean(self, probs): + # probs: [..., num_classes] -> [..., 2] (E[cont | probs]) + return probs @ self.action_table.to(probs.dtype) diff --git a/pufferlib/pufferl.py b/pufferlib/pufferl.py index 4dabbc21a9..c8456d9c12 100644 --- a/pufferlib/pufferl.py +++ b/pufferlib/pufferl.py @@ -404,8 +404,7 @@ def evaluate(self): logits, value = self.policy.forward_eval(o_device, state) logits = logits_to_float(logits) - action, logprob, _, cont_action = pufferlib.pytorch.sample_logits(logits, action=None, deterministic=False, - env_continuous=self.env_continuous, policy=self.policy) + action, logprob, _, cont_action = pufferlib.pytorch.sample_logits(logits, env_continuous=self.env_continuous, policy=self.policy) if config["normalize_rewards"]: r = torch.sign(r) * torch.log1p(torch.abs(r)) @@ -1855,7 +1854,7 @@ def mine_failures(env_name, args=None): o_t = torch.as_tensor(obs_arr).to(device) state = {"reward": None, "done": None, "env_id": None, "mask": None} logits, _ = policy.forward_eval(o_t, state) - action, _, _, cont_action = pufferlib.pytorch.sample_logits(logits, action=None, deterministic=False, env_continuous=env_continuous, policy=policy) + action, _, _, cont_action = pufferlib.pytorch.sample_logits(logits, env_continuous=env_continuous, policy=policy) action = action.cpu().numpy() if action.ndim == 1 and len(vecenv.single_action_space.shape) >= 1: action = action.reshape(-1, *vecenv.single_action_space.shape) diff --git a/pufferlib/pytorch.py b/pufferlib/pytorch.py index 466b1126ec..871d8a192e 100644 --- a/pufferlib/pytorch.py +++ b/pufferlib/pytorch.py @@ -20,6 +20,9 @@ np.dtype("int8"): torch.int8, } +ACTION_SELECT_SAMPLE = "sample" +ACTION_SELECT_MODE = "mode" # argmax +ACTION_SELECT_MEAN = "mean" # probability-weighted continuous mean LITTLE_BYTE_ORDER = sys.byteorder == "little" @@ -187,13 +190,15 @@ def entropy_probs(logits, probs): return -p_log_p.sum(-1) -def sample_logits(logits, action=None, deterministic=False, env_continuous=None, policy=None): # TODO discrete continuous +def sample_logits(logits, action=None, action_selection=ACTION_SELECT_SAMPLE, env_continuous=None, policy=None): # TODO discrete continuous is_discrete = isinstance(logits, torch.Tensor) + if action_selection == ACTION_SELECT_MEAN and not (env_continuous and not policy.is_continuous): + raise ValueError("action_selection='mean' requires a discrete policy on a continuous env") if isinstance(logits, torch.distributions.Normal): batch = logits.loc.shape[0] if action is None: action = logits.sample().view(batch, -1) - if deterministic: + if action_selection != ACTION_SELECT_SAMPLE: action = logits.loc.view(batch, -1) # TODO - DETERMINISTIC use mean action for eval log_probs = logits.log_prob(action.view(batch, -1)).sum(1) @@ -211,14 +216,13 @@ def sample_logits(logits, action=None, deterministic=False, env_continuous=None, normalized_logits = logits - logits.logsumexp(dim=-1, keepdim=True) probs = logits_to_probs(logits) - # TODO this jumps in here and makes the action determinisitc if action is None: - if deterministic: - action = torch.argmax(probs, -1) # TODO weighted average - else: + if action_selection == ACTION_SELECT_SAMPLE: probs = torch.nan_to_num(probs, 1e-8, 1e-8, 1e-8) action = torch.multinomial(probs.reshape(-1, probs.shape[-1]), 1, replacement=True).int() action = action.reshape(probs.shape[:-1]) + else: # MODE and MEAN both use argmax as the *nominal* discrete action (for logging) + action = torch.argmax(probs, -1) else: batch = logits[0].shape[0] action = action.view(batch, -1).T @@ -228,11 +232,12 @@ def sample_logits(logits, action=None, deterministic=False, env_continuous=None, logits_entropy = entropy(normalized_logits).sum(0) if env_continuous and not policy.is_continuous: - cont_actions = policy.discrete_actions_to_continuous(action) - + if action_selection == ACTION_SELECT_MEAN: + cont_actions = policy.discrete_probs_to_continuous_mean(probs) + else: # Mode and sample we already have an action from the code above. + cont_actions = policy.discrete_actions_to_continuous(action) if is_discrete: return action.squeeze(0), logprob.squeeze(0), logits_entropy.squeeze(0), cont_actions.squeeze(0) - return action.T, logprob.sum(0), logits_entropy, cont_actions if is_discrete: From 96662c1a32112e72ddfbed01451a016015dcf75f Mon Sep 17 00:00:00 2001 From: Bernhard Jaeger Date: Wed, 22 Jul 2026 19:07:19 +0200 Subject: [PATCH 19/35] Fix eval bug. --- pufferlib/ocean/benchmark/evaluators/base.py | 21 +++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/pufferlib/ocean/benchmark/evaluators/base.py b/pufferlib/ocean/benchmark/evaluators/base.py index eaa8cbe266..487421566d 100644 --- a/pufferlib/ocean/benchmark/evaluators/base.py +++ b/pufferlib/ocean/benchmark/evaluators/base.py @@ -689,7 +689,6 @@ def _render_pass_obs(self, vecenv, policy, args) -> list: logits, value = policy.forward_eval(ob_t, state) pool_outputs = pool_method(ob_t, state) if pool_method is not None else {} action, logprob, entropy, cont_action = pufferlib.pytorch.sample_logits(logits, action_selection=pufferlib.pytorch.ACTION_SELECT_MODE, env_continuous=env_continuous, policy=policy) - raw_action = action.cpu().numpy().reshape(vec.action_space.shape) pool_outputs = {k: v.cpu().numpy().astype(np.int16, copy=False) for k, v in pool_outputs.items()} if pool_hist is None and pool_outputs: pool_hist = { @@ -699,9 +698,19 @@ def _render_pass_obs(self, vecenv, policy, args) -> list: ] for k, values in pool_outputs.items() } - clipped_action = raw_action + + if env_continuous and not policy.is_continuous: # TODO check with trajectory + raw_action = cont_action.cpu().numpy().reshape(vec.action_space.shape) + clipped_action = raw_action + ob, _, _, _, step_infos = vec.step(clipped_action) + else: + raw_action = action.cpu().numpy().reshape(vec.action_space.shape) + clipped_action = raw_action + if isinstance(logits, torch.distributions.Normal): + clipped_action = np.clip(raw_action, vec.action_space.low, vec.action_space.high) + ob, _, _, _, step_infos = vec.step(clipped_action) + if isinstance(logits, torch.distributions.Normal): - clipped_action = np.clip(raw_action, vec.action_space.low, vec.action_space.high) policy_outputs = { "mean": logits.loc.cpu().numpy().reshape(vec.action_space.shape), "std": logits.scale.cpu().numpy().reshape(vec.action_space.shape), @@ -759,12 +768,6 @@ def _render_pass_obs(self, vecenv, policy, args) -> list: np.asarray(policy_outputs[start_obs_index:end_obs_index], dtype=np.float32).copy() ) start_obs_index = end_obs_index - - if env_continuous and not policy.is_continuous: # TODO check with trajectory - cont_action = cont_action.cpu().numpy().reshape(vec.action_space.shape) - ob, _, _, _, step_infos = vec.step(cont_action) - else: - ob, _, _, _, step_infos = vec.step(clipped_action) for d in self._flatten_infos(step_infos): if isinstance(d, dict) and d.get("summary_type") == "completed_episode": From 27fd5df0d58cbe20dc4a9d48043da08d1b80ec9f Mon Sep 17 00:00:00 2001 From: Bernhard Jaeger Date: Wed, 22 Jul 2026 19:37:00 +0200 Subject: [PATCH 20/35] Made deterministic option a settable paramter. --- pufferlib/config/puffer_drive.yaml | 8 +++++++ pufferlib/ocean/benchmark/evaluators/base.py | 25 ++++++++++++++++---- 2 files changed, 29 insertions(+), 4 deletions(-) diff --git a/pufferlib/config/puffer_drive.yaml b/pufferlib/config/puffer_drive.yaml index 22d0eef9c4..d8cebb25f9 100644 --- a/pufferlib/config/puffer_drive.yaml +++ b/pufferlib/config/puffer_drive.yaml @@ -333,6 +333,10 @@ eval: obs_dropout_boundary: 0.0 obs_slots_lane_n: 80 obs_slots_boundary_n: 80 + # Policy action selection during eval: sample | mode | mean. + # mode = argmax; mean = probability-weighted continuous action + # (mean requires a discrete policy on a continuous env). + action_selection: mode eval: num_scenarios: 250 export_episode_csv: 'true' @@ -409,6 +413,8 @@ eval: init_mode: create_all_valid scenario_length: 201 obs_slots_partners_n: 32 + # Policy action selection during eval: sample | mode | mean. + action_selection: mode eval: num_scenarios: 50 render_num_scenarios: 2 @@ -500,6 +506,8 @@ eval: init_mode: create_all_valid init_step: 10 goal_radius: 2.0 + # Policy action selection during eval: sample | mode | mean. + action_selection: mode eval: wosac_num_rollouts: 32 wosac_num_agents: 256 diff --git a/pufferlib/ocean/benchmark/evaluators/base.py b/pufferlib/ocean/benchmark/evaluators/base.py index 487421566d..c7b38896d4 100644 --- a/pufferlib/ocean/benchmark/evaluators/base.py +++ b/pufferlib/ocean/benchmark/evaluators/base.py @@ -71,6 +71,23 @@ def __init__(self, name: str, config: dict, train_config: dict): self.render_views: list = list(config.get("render_views", ["sim_state"])) self.clean: bool = bool(config.get("clean", True)) + # Policy action selection during eval: sample | mode | mean. A top-level + # evaluator knob (like enabled/render), distinct from `self.mode` + # (inline/subprocess run mode). Defaults to argmax (mode). + import pufferlib.pytorch + + self.action_selection: str = config.get("action_selection", pufferlib.pytorch.ACTION_SELECT_MODE) + valid_action_selections = ( + pufferlib.pytorch.ACTION_SELECT_SAMPLE, + pufferlib.pytorch.ACTION_SELECT_MODE, + pufferlib.pytorch.ACTION_SELECT_MEAN, + ) + if self.action_selection not in valid_action_selections: + raise ValueError( + f"[eval.{name}].action_selection='{self.action_selection}' " + f"must be one of {valid_action_selections}" + ) + # -- Config hooks --------------------------------------------------- def env_overrides(self) -> dict: @@ -155,7 +172,7 @@ def _run_rollout_loop(self, vecenv, policy, args) -> dict: with torch.no_grad(): ob_t = torch.as_tensor(obs).to(device) logits, _ = policy.forward_eval(ob_t, state) - action, _, _, cont_action = pufferlib.pytorch.sample_logits(logits, action_selection=pufferlib.pytorch.ACTION_SELECT_MODE, env_continuous=env_continuous, policy=policy) + action, _, _, cont_action = pufferlib.pytorch.sample_logits(logits, action_selection=self.action_selection, env_continuous=env_continuous, policy=policy) if env_continuous and not policy.is_continuous: # TODO check with trajectory cont_action = cont_action.cpu().numpy().reshape(vecenv.action_space.shape) @@ -523,7 +540,7 @@ def _render_pass_html(self, vecenv, policy, args) -> list: with torch.no_grad(): ob_t = torch.as_tensor(ob).to(device) logits, _ = policy.forward_eval(ob_t, state) - action, _, _, cont_action = pufferlib.pytorch.sample_logits(logits, action_selection=pufferlib.pytorch.ACTION_SELECT_MODE, env_continuous=env_continuous, policy=policy) + action, _, _, cont_action = pufferlib.pytorch.sample_logits(logits, action_selection=self.action_selection, env_continuous=env_continuous, policy=policy) if env_continuous and not policy.is_continuous: # TODO check with trajectory cont_action = cont_action.cpu().numpy().reshape(vecenv.action_space.shape) @@ -688,7 +705,7 @@ def _render_pass_obs(self, vecenv, policy, args) -> list: ob_t = torch.as_tensor(ob).to(device) logits, value = policy.forward_eval(ob_t, state) pool_outputs = pool_method(ob_t, state) if pool_method is not None else {} - action, logprob, entropy, cont_action = pufferlib.pytorch.sample_logits(logits, action_selection=pufferlib.pytorch.ACTION_SELECT_MODE, env_continuous=env_continuous, policy=policy) + action, logprob, entropy, cont_action = pufferlib.pytorch.sample_logits(logits, action_selection=self.action_selection, env_continuous=env_continuous, policy=policy) pool_outputs = {k: v.cpu().numpy().astype(np.int16, copy=False) for k, v in pool_outputs.items()} if pool_hist is None and pool_outputs: pool_hist = { @@ -877,7 +894,7 @@ def _render_view(self, vecenv, target_env, policy, args, view_idx, out_dir, view with torch.no_grad(): ob_t = torch.as_tensor(ob).to(device) logits, _ = policy.forward_eval(ob_t, state) - action, _, _, cont_action = pufferlib.pytorch.sample_logits(logits, action_selection=pufferlib.pytorch.ACTION_SELECT_MODE, env_continuous=env_continuous, policy=policy) + action, _, _, cont_action = pufferlib.pytorch.sample_logits(logits, action_selection=self.action_selection, env_continuous=env_continuous, policy=policy) if env_continuous and not policy.is_continuous: # TODO check with trajectory cont_action = cont_action.cpu().numpy().reshape(vecenv.action_space.shape) From da028246231edf811fa617109514337c87695f92 Mon Sep 17 00:00:00 2001 From: bernhard_jaeger Date: Thu, 23 Jul 2026 18:12:02 +0000 Subject: [PATCH 21/35] remove broken evals --- pufferlib/config/ocean/drive.ini | 88 ----------------------- pufferlib/config/puffer_drive.yaml | 112 +---------------------------- scripts/1_train_model_slurm.sh | 8 +-- 3 files changed, 5 insertions(+), 203 deletions(-) diff --git a/pufferlib/config/ocean/drive.ini b/pufferlib/config/ocean/drive.ini index 8501127fc8..6c53556c7c 100644 --- a/pufferlib/config/ocean/drive.ini +++ b/pufferlib/config/ocean/drive.ini @@ -299,22 +299,6 @@ eval.num_scenarios = 250 eval.export_episode_csv = true eval.verify_coverage = true -; [eval.validation_replay] -; inherits = "validation_defaults" -; type = "multi_scenario" -; enabled = true -; render = true -; render_backend = "triage_html" -; env.simulation_mode = "replay" -; env.control_mode = "control_sdc_only" -; env.map_dir = "/lustre/scwpod02/client/kyutai/kesai/data/nuPlan/PufferDrive" -; env.num_maps = 250 -; env.max_agents_per_env = 64 -; env.scenario_length = 200 -; env.resample_frequency = 200 -; eval.render_num_scenarios = 5 -; eval.render_max_steps = 200 - [eval.validation_gigaflow] inherits = "validation_defaults" type = "multi_scenario" @@ -372,78 +356,6 @@ eval.num_scenarios = 50 eval.render_num_scenarios = 2 eval.render_max_steps = 200 -; [eval.behaviors_full_dir] -; inherits = "behaviors_defaults" -; type = "behavior_class" -; enabled = true -; env.map_dir = "/lustre/scwpod02/client/kyutai/kesai/data/nuPlan/PufferDrive" - -; [eval.behaviors_hard_stop] -; inherits = "behaviors_defaults" -; type = "behavior_class" -; enabled = true -; env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/hard_stop" -; -; [eval.behaviors_highway_straight] -; inherits = "behaviors_defaults" -; type = "behavior_class" -; enabled = true -; env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/highway_straight" -; -; [eval.behaviors_lane_change] -; inherits = "behaviors_defaults" -; type = "behavior_class" -; enabled = true -; env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/lane_change" -; -; [eval.behaviors_merge] -; inherits = "behaviors_defaults" -; type = "behavior_class" -; enabled = true -; env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/merge" -; -; [eval.behaviors_parked_cars] -; inherits = "behaviors_defaults" -; type = "behavior_class" -; enabled = true -; env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/parked_cars" -; -; [eval.behaviors_roundabout] -; inherits = "behaviors_defaults" -; type = "behavior_class" -; enabled = true -; env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/roundabout" -; -; [eval.behaviors_stopped_traffic] -; inherits = "behaviors_defaults" -; type = "behavior_class" -; enabled = true -; env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/stopped_traffic" -; -; [eval.behaviors_traffic_light_green] -; inherits = "behaviors_defaults" -; type = "behavior_class" -; enabled = true -; env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/traffic_light_green" -; -; [eval.behaviors_traffic_light_stop] -; inherits = "behaviors_defaults" -; type = "behavior_class" -; enabled = true -; env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/traffic_light_stop" -; -; [eval.behaviors_unprotected_left] -; inherits = "behaviors_defaults" -; type = "behavior_class" -; enabled = true -; env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/unprotected_left" -; -; [eval.behaviors_unprotected_right] -; inherits = "behaviors_defaults" -; type = "behavior_class" -; enabled = true -; env.map_dir = "/scratch/ev2237/data/nuplan/categories_v021/unprotected_right" - ; --------------------------------------------------------------------------- ; Optional: WOSAC realism eval. Off by default. ; --------------------------------------------------------------------------- diff --git a/pufferlib/config/puffer_drive.yaml b/pufferlib/config/puffer_drive.yaml index 54da056a0f..453c18aad2 100644 --- a/pufferlib/config/puffer_drive.yaml +++ b/pufferlib/config/puffer_drive.yaml @@ -336,23 +336,6 @@ eval: num_scenarios: 250 export_episode_csv: 'true' verify_coverage: 'true' - validation_replay: - inherits: validation_defaults - type: multi_scenario - enabled: 'true' - render: 'true' - render_backend: triage_html - env: - simulation_mode: replay - control_mode: control_sdc_only - map_dir: /scratch/ev2237/data/nuplan/nuplan_mini_train_bins - num_maps: 250 - max_agents_per_env: 64 - scenario_length: 200 - resample_frequency: 200 - eval: - render_num_scenarios: 5 - render_max_steps: 200 validation_gigaflow: inherits: validation_defaults type: multi_scenario @@ -390,100 +373,7 @@ eval: num_scenarios: 32 render_num_scenarios: 16 render_max_steps: 300 - # --------------------------------------------------------------------------- - # Driving-behaviour evaluation: nuPlan scenes labeled by scene type. Each - # behavior is one [eval.behaviors_*] section. All inherit from the template - # below — change shared knobs in one place. - # --------------------------------------------------------------------------- - behaviors_defaults: - enabled: 'false' - interval: 250 - mode: inline - clean: 'true' - render: 'true' - render_views: [sim_state, bev] - env: - simulation_mode: replay - control_mode: control_sdc_only - init_mode: create_all_valid - scenario_length: 201 - obs_slots_partners_n: 32 - eval: - num_scenarios: 50 - render_num_scenarios: 2 - render_max_steps: 200 - behaviors_full_dir: - inherits: behaviors_defaults - type: behavior_class - enabled: 'true' - env: - map_dir: /scratch/ev2237/data/nuplan/nuplan_mini_train_bins - behaviors_hard_stop: - inherits: behaviors_defaults - type: behavior_class - enabled: 'true' - env: - map_dir: /scratch/ev2237/data/nuplan/categories_v021/hard_stop - behaviors_highway_straight: - inherits: behaviors_defaults - type: behavior_class - enabled: 'true' - env: - map_dir: /scratch/ev2237/data/nuplan/categories_v021/highway_straight - behaviors_lane_change: - inherits: behaviors_defaults - type: behavior_class - enabled: 'true' - env: - map_dir: /scratch/ev2237/data/nuplan/categories_v021/lane_change - behaviors_merge: - inherits: behaviors_defaults - type: behavior_class - enabled: 'true' - env: - map_dir: /scratch/ev2237/data/nuplan/categories_v021/merge - behaviors_parked_cars: - inherits: behaviors_defaults - type: behavior_class - enabled: 'true' - env: - map_dir: /scratch/ev2237/data/nuplan/categories_v021/parked_cars - behaviors_roundabout: - inherits: behaviors_defaults - type: behavior_class - enabled: 'true' - env: - map_dir: /scratch/ev2237/data/nuplan/categories_v021/roundabout - behaviors_stopped_traffic: - inherits: behaviors_defaults - type: behavior_class - enabled: 'true' - env: - map_dir: /scratch/ev2237/data/nuplan/categories_v021/stopped_traffic - behaviors_traffic_light_green: - inherits: behaviors_defaults - type: behavior_class - enabled: 'true' - env: - map_dir: /scratch/ev2237/data/nuplan/categories_v021/traffic_light_green - behaviors_traffic_light_stop: - inherits: behaviors_defaults - type: behavior_class - enabled: 'true' - env: - map_dir: /scratch/ev2237/data/nuplan/categories_v021/traffic_light_stop - behaviors_unprotected_left: - inherits: behaviors_defaults - type: behavior_class - enabled: 'true' - env: - map_dir: /scratch/ev2237/data/nuplan/categories_v021/unprotected_left - behaviors_unprotected_right: - inherits: behaviors_defaults - type: behavior_class - enabled: 'true' - env: - map_dir: /scratch/ev2237/data/nuplan/categories_v021/unprotected_right + # --------------------------------------------------------------------------- # Optional: WOSAC realism eval. Off by default. # --------------------------------------------------------------------------- diff --git a/scripts/1_train_model_slurm.sh b/scripts/1_train_model_slurm.sh index f00da483c6..0108810df6 100644 --- a/scripts/1_train_model_slurm.sh +++ b/scripts/1_train_model_slurm.sh @@ -6,15 +6,15 @@ #SBATCH --gres gpu:8 #SBATCH --mem=1007G #SBATCH --cpus-per-task 144 -#SBATCH --output /home/bjaeger/PufferDrive/experiments/k_nightly_0001/log_%a_%A.out -#SBATCH --error /home/bjaeger/PufferDrive/experiments/k_nightly_0001/log_%a_%A.err +#SBATCH --output /home/bjaeger/PufferDrive/experiments/k_nightly_0005/log_%a_%A.out +#SBATCH --error /home/bjaeger/PufferDrive/experiments/k_nightly_0005/log_%a_%A.err #SBATCH --partition dev # print info about current job echo "START TIME: $(date)" start=$(date +%s) -export RUN_NAME=k_nightly_0001 +export RUN_NAME=k_nightly_0005 echo ${RUN_NAME} # TODO could try to tune these. 1 Is probably best since Puffer parallelizes across all cores. @@ -40,7 +40,7 @@ torchrun --standalone --nnodes=1 --nproc-per-node=8 --max_restarts=0 --start-met train.max_minibatch_size=131072 \ train.minibatch_size=131072 \ train.precision=bfloat16 \ - env.num_agents=4096 \ + eval.validation_gigaflow.render_backend=obs_html \ tb=True From 1558636c8e8fdeddbf64b50cb7e6ffcdf10d03e6 Mon Sep 17 00:00:00 2001 From: bernhard_jaeger Date: Thu, 23 Jul 2026 18:12:32 +0000 Subject: [PATCH 22/35] remove broken evals --- scripts/2_train_model_multinode_slurm.sh | 63 ++++++++++++++++++++++++ 1 file changed, 63 insertions(+) create mode 100644 scripts/2_train_model_multinode_slurm.sh diff --git a/scripts/2_train_model_multinode_slurm.sh b/scripts/2_train_model_multinode_slurm.sh new file mode 100644 index 0000000000..18d7940bae --- /dev/null +++ b/scripts/2_train_model_multinode_slurm.sh @@ -0,0 +1,63 @@ +#!/bin/bash +#SBATCH --job-name train_puffer +#SBATCH --nodes 16 # Number of nodes requested +#SBATCH --ntasks-per-node 1 # Run 1 srun task per node (which fires up torchrun) +#SBATCH --gres gpu:8 # GPUs per node +#SBATCH --cpus-per-task 144 +#SBATCH --mem=1007G +#SBATCH --time 2-00:00 +#SBATCH --output /home/bjaeger/PufferDrive/experiments/k_nightly_0003/log_%a_%A.out +#SBATCH --error /home/bjaeger/PufferDrive/experiments/k_nightly_0003/log_%a_%A.err +#SBATCH --partition dev + +# Set up PyTorch Distributed Rendezvous parameters from Slurm variables +export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1) +export MASTER_PORT=29500 +export WORLD_SIZE=$SLURM_JOB_NUM_NODES + +echo "START TIME: $(date)" +echo "Master node: ${MASTER_ADDR}:${MASTER_PORT}" +echo "Total nodes: ${WORLD_SIZE}" +start=$(date +%s) + +export RUN_NAME=k_nightly_0003 +echo ${RUN_NAME} + +# Thread limit limits CPU thrashing across worker environments +export NUMEXPR_NUM_THREADS=1 +export MKL_NUM_THREADS=1 +export VECLIB_MAXIMUM_THREADS=1 +export OPENBLAS_NUM_THREADS=1 +export OMP_NUM_THREADS=1 + +source .venv/bin/activate + +# Execute torchrun across all nodes using srun +srun torchrun \ + --nnodes=${SLURM_JOB_NUM_NODES} \ + --nproc-per-node=8 \ + --rdzv_id=${SLURM_JOB_ID} \ + --rdzv_backend=c10d \ + --rdzv_endpoint="${MASTER_ADDR}:${MASTER_PORT}" \ + --max_restarts=0 \ + --start-method spawn \ + -m pufferlib.pufferl train puffer_drive \ + wandb=True \ + wandb_project=nightly-multi-long \ + wandb_group=emerge_ \ + train.data_dir=/home/bjaeger/PufferDrive/experiments/${RUN_NAME} \ + env.map_dir=/home/bjaeger/PufferDrive/pufferlib/resources/drive/binaries/carla \ + train.name=${RUN_NAME} \ + train.total_timesteps=100000000000 \ + vec.num_envs=16 \ + +eval.map_dir=/home/bjaeger/data/nuPlan/PufferDrive \ + train.compile=True \ + train.max_minibatch_size=131072 \ + train.minibatch_size=131072 \ + train.precision=bfloat16 \ + tb=True + +end=$(date +%s) +runtime=$((end-start)) +echo "END TIME: $(date)" +echo "Runtime: ${runtime}" From 9a8faceb938d8d01c0061b086026ee6bce6f2f53 Mon Sep 17 00:00:00 2001 From: Bernhard Jaeger Date: Thu, 23 Jul 2026 20:36:33 +0200 Subject: [PATCH 23/35] Changed default env action to continuous as the model now does the discretization. --- pufferlib/config/ocean/drive.ini | 2 +- pufferlib/config/puffer_drive.yaml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pufferlib/config/ocean/drive.ini b/pufferlib/config/ocean/drive.ini index 134494ffa3..c5e82cb0af 100644 --- a/pufferlib/config/ocean/drive.ini +++ b/pufferlib/config/ocean/drive.ini @@ -53,7 +53,7 @@ num_agents = 1024 min_agents_per_env = 1 max_agents_per_env = 80 ; Actions type - options: "discrete", "continuous" -action_type = "discrete" +action_type = "continuous" ; Dynamics model - options: "classic", "jerk" dynamics_model = "jerk" ; Time delta between steps in seconds diff --git a/pufferlib/config/puffer_drive.yaml b/pufferlib/config/puffer_drive.yaml index 81a4842157..9c2f97206b 100644 --- a/pufferlib/config/puffer_drive.yaml +++ b/pufferlib/config/puffer_drive.yaml @@ -47,7 +47,7 @@ env: min_agents_per_env: 1 max_agents_per_env: 80 # Actions type - options: "discrete", "continuous" - action_type: discrete + action_type: continuous # Dynamics model - options: "classic", "jerk" dynamics_model: jerk # Time delta between steps in seconds From d6ea46685c7e1aef93d655a968a442570dd54523 Mon Sep 17 00:00:00 2001 From: Bernhard Jaeger Date: Thu, 23 Jul 2026 20:51:29 +0200 Subject: [PATCH 24/35] Train next nightly --- scripts/1_train_model_slurm.sh | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/scripts/1_train_model_slurm.sh b/scripts/1_train_model_slurm.sh index 0108810df6..4c9fe9f8bc 100644 --- a/scripts/1_train_model_slurm.sh +++ b/scripts/1_train_model_slurm.sh @@ -6,15 +6,15 @@ #SBATCH --gres gpu:8 #SBATCH --mem=1007G #SBATCH --cpus-per-task 144 -#SBATCH --output /home/bjaeger/PufferDrive/experiments/k_nightly_0005/log_%a_%A.out -#SBATCH --error /home/bjaeger/PufferDrive/experiments/k_nightly_0005/log_%a_%A.err +#SBATCH --output /home/bjaeger/PufferDrive/experiments/k_nightly_0006/log_%a_%A.out +#SBATCH --error /home/bjaeger/PufferDrive/experiments/k_nightly_0006/log_%a_%A.err #SBATCH --partition dev # print info about current job echo "START TIME: $(date)" start=$(date +%s) -export RUN_NAME=k_nightly_0005 +export RUN_NAME=k_nightly_0006 echo ${RUN_NAME} # TODO could try to tune these. 1 Is probably best since Puffer parallelizes across all cores. From 7a711b1e9446b2ff93d915eb61762c10ac92bff0 Mon Sep 17 00:00:00 2001 From: Bernhard Jaeger Date: Thu, 23 Jul 2026 20:57:05 +0200 Subject: [PATCH 25/35] Run autoformatter to make the git pull request happy. --- bernhard_configs.txt | 4 +-- bernhard_notes.txt | 16 ++++----- notebooks/05_inference.py | 4 ++- notebooks/06_architecture.py | 2 +- notebooks/notebook_utils.py | 2 +- pufferlib/config/puffer_drive.yaml | 4 +-- pufferlib/ocean/benchmark/evaluator.py | 19 ++++++----- pufferlib/ocean/benchmark/evaluators/base.py | 31 ++++++++++------- pufferlib/ocean/torch.py | 36 ++++++++++++++++---- pufferlib/pufferl.py | 13 ++++--- pufferlib/pytorch.py | 10 +++--- scripts/0_build_puffer.sh | 2 +- scripts/0_eval_model_slurm.sh | 2 +- scripts/0_train_model_slurm.sh | 2 +- scripts/cluster_configs/bernhard_1.yaml | 2 +- scripts/cluster_configs/bernhard_8.yaml | 2 +- 16 files changed, 94 insertions(+), 57 deletions(-) diff --git a/bernhard_configs.txt b/bernhard_configs.txt index 682b7c5d1e..5920f9a1fe 100644 --- a/bernhard_configs.txt +++ b/bernhard_configs.txt @@ -1,6 +1,6 @@ puffer train puffer_drive --vec.num-workers 1 --vec.num-envs 1 --env.num-agents 32 --train.minibatch-size 256 -puffer train puffer_drive --train.learning_rate 0.001 +puffer train puffer_drive --train.learning_rate 0.001 .venv/bin/puffer train puffer_drive --wandb --vec.num-envs 16 --train.precision bfloat16 --train.compile 1 @@ -331,4 +331,4 @@ Usage=puffer [--load-model-path LOAD_MODEL_PATH] [--load-id LOAD_ID] [--render-m [--eval.wosac.eval.wosac-sanity-check EVAL.WOSAC.EVAL.WOSAC_SANITY_CHECK] [--eval.wosac.eval.wosac-aggregate-results EVAL.WOSAC.EVAL.WOSAC_AGGREGATE_RESULTS] [--controlled-exp.train.learning-rate.values CONTROLLED_EXP.TRAIN.LEARNING_RATE.VALUES] - [--controlled-exp.train.ent-coef.values CONTROLLED_EXP.TRAIN.ENT_COEF.VALUES] [-h] \ No newline at end of file + [--controlled-exp.train.ent-coef.values CONTROLLED_EXP.TRAIN.ENT_COEF.VALUES] [-h] diff --git a/bernhard_notes.txt b/bernhard_notes.txt index 0e3ea21dde..765a5e3205 100644 --- a/bernhard_notes.txt +++ b/bernhard_notes.txt @@ -2,7 +2,7 @@ git merge origin/3.0 git push origin bernhard -# Not in valentins config but in my config. +# Not in valentins config but in my config. # non_sdc_controller: policy # non_vehicle_controller: auto # sdc_controller: policy @@ -48,11 +48,11 @@ SPS 341.4 metrics = self._run_rollout_loop(vecenv, policy, args) File "/lustre/scwpod02/client/kyutai/kesai/bernhard/PufferDrive/experiments/bernhard_train_base_wandb1_run-name003_longrun_total-timesteps300000000000_num-envs24_max-waypoint-spacing40.0_num-agents2048_obs-norm-ve_067a078/code_v1/pufferlib/ocean/benchmark/evaluators/base.py", line 156, in _run_rollout_loop logits, _ = policy.forward_eval(ob_t, state) - - - - - + + + + + slide_idx 687 traffic_control_dim 28 self.obs_valid_count_features 4 @@ -114,7 +114,7 @@ lauch.json: "--vec.num-envs", "16", "--env.num-agents", "1024", "--train.minibatch-size", "131072", - "--train.max-minibatch-size", "131072", + "--train.max-minibatch-size", "131072", ], "console"="integratedTerminal", "justMyCode"=false @@ -215,4 +215,4 @@ With Expo and PBO │ reward_components/velocity 0.802 reward_components/timestep -0.000 │ │ reward_components/reverse -0.000 reward_components/overspeed -0.003 │ │ reward_components/ade 0.000 │ -╰───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ \ No newline at end of file +╰───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ diff --git a/notebooks/05_inference.py b/notebooks/05_inference.py index 8e5538514b..187f4f4d58 100644 --- a/notebooks/05_inference.py +++ b/notebooks/05_inference.py @@ -79,7 +79,9 @@ # Sample actions action, logprob, ent, cont_action = sample_logits(logits_list, env_continuous=env_continuous, policy=policy) -action_det, _, _, _ = sample_logits(logits_list, action_selection=ACTION_SELECT_MODE, env_continuous=env_continuous, policy=policy) +action_det, _, _, _ = sample_logits( + logits_list, action_selection=ACTION_SELECT_MODE, env_continuous=env_continuous, policy=policy +) print(f"Value: mean={value.mean():.4f}, std={value.std():.4f}, range=[{value.min():.4f}, {value.max():.4f}]") print(f"Entropy: mean={ent.mean():.4f}, std={ent.std():.4f}") diff --git a/notebooks/06_architecture.py b/notebooks/06_architecture.py index 75171786e0..671ff33c33 100644 --- a/notebooks/06_architecture.py +++ b/notebooks/06_architecture.py @@ -521,7 +521,7 @@ def count_params(module): "backbone_layer_norm": False, "shared_network": True, "mask_padded_features": False, - "action_type": "discrete" + "action_type": "discrete", } results = [] diff --git a/notebooks/notebook_utils.py b/notebooks/notebook_utils.py index 3b0c05b4e0..8feb99bea1 100644 --- a/notebooks/notebook_utils.py +++ b/notebooks/notebook_utils.py @@ -112,7 +112,7 @@ "backbone_layer_norm": False, "shared_network": True, "mask_padded_features": False, - "action_type": "discrete" + "action_type": "discrete", } diff --git a/pufferlib/config/puffer_drive.yaml b/pufferlib/config/puffer_drive.yaml index 9c2f97206b..7b9b43f911 100644 --- a/pufferlib/config/puffer_drive.yaml +++ b/pufferlib/config/puffer_drive.yaml @@ -378,8 +378,8 @@ eval: num_scenarios: 32 render_num_scenarios: 16 render_max_steps: 300 - - + + # --------------------------------------------------------------------------- # Driving-behaviour evaluation: nuPlan scenes labeled by scene type. Each # behavior is one [eval.behaviors_*] section. All inherit from the template diff --git a/pufferlib/ocean/benchmark/evaluator.py b/pufferlib/ocean/benchmark/evaluator.py index e42927c5b7..133b37f24b 100644 --- a/pufferlib/ocean/benchmark/evaluator.py +++ b/pufferlib/ocean/benchmark/evaluator.py @@ -88,7 +88,7 @@ def collect_simulated_trajectories(self, args, puffer_env, policy): "id": np.zeros((num_agents, self.num_rollouts, self.sim_steps), dtype=np.int32), } - env_continuous = isinstance(puffer_env.single_action_space, pufferlib.spaces.Box) + env_continuous = isinstance(puffer_env.single_action_space, pufferlib.spaces.Box) for rollout_idx in range(self.num_rollouts): print(f"\rCollecting rollout {rollout_idx + 1}/{self.num_rollouts}...", end="", flush=True) @@ -114,19 +114,19 @@ def collect_simulated_trajectories(self, args, puffer_env, policy): with torch.no_grad(): ob_tensor = torch.as_tensor(obs).to(device) logits, value = policy.forward_eval(ob_tensor, state) - action, logprob, _, cont_action = pufferlib.pytorch.sample_logits(logits, env_continuous=env_continuous, policy=policy) + action, logprob, _, cont_action = pufferlib.pytorch.sample_logits( + logits, env_continuous=env_continuous, policy=policy + ) action_np = action.cpu().numpy().reshape(puffer_env.action_space.shape) if isinstance(logits, torch.distributions.Normal): action_np = np.clip(action_np, puffer_env.action_space.low, puffer_env.action_space.high) - - if env_continuous and not policy.is_continuous: # TODO check with trajectory + if env_continuous and not policy.is_continuous: # TODO check with trajectory cont_action = cont_action.cpu().numpy().reshape(puffer_env.action_space.shape) obs, _, _, _, _ = puffer_env.step(cont_action) else: obs, _, _, _, _ = puffer_env.step(action_np) - return trajectories @@ -668,7 +668,7 @@ def rollout(self, args, puffer_env, policy): total_steps = (scenario_length - init_steps + 1) * num_maps obs, info = puffer_env.reset() - env_continuous = isinstance(puffer_env.single_action_space, pufferlib.spaces.Box) + env_continuous = isinstance(puffer_env.single_action_space, pufferlib.spaces.Box) state = {} if args["train"]["use_rnn"]: state = dict( @@ -681,14 +681,15 @@ def rollout(self, args, puffer_env, policy): with torch.no_grad(): ob_tensor = torch.as_tensor(obs).to(device) logits, value = policy.forward_eval(ob_tensor, state) - action, logprob, _, cont_action = pufferlib.pytorch.sample_logits(logits, env_continuous=env_continuous, policy=policy) + action, logprob, _, cont_action = pufferlib.pytorch.sample_logits( + logits, env_continuous=env_continuous, policy=policy + ) action_np = action.cpu().numpy().reshape(puffer_env.action_space.shape) if isinstance(logits, torch.distributions.Normal): action_np = np.clip(action_np, puffer_env.action_space.low, puffer_env.action_space.high) - - if env_continuous and not policy.is_continuous: # TODO check with trajectory + if env_continuous and not policy.is_continuous: # TODO check with trajectory cont_action = cont_action.cpu().numpy().reshape(puffer_env.action_space.shape) obs, rewards, dones, truncs, info_list = puffer_env.step(cont_action) else: diff --git a/pufferlib/ocean/benchmark/evaluators/base.py b/pufferlib/ocean/benchmark/evaluators/base.py index c7b38896d4..61c2ed9bc2 100644 --- a/pufferlib/ocean/benchmark/evaluators/base.py +++ b/pufferlib/ocean/benchmark/evaluators/base.py @@ -84,8 +84,7 @@ def __init__(self, name: str, config: dict, train_config: dict): ) if self.action_selection not in valid_action_selections: raise ValueError( - f"[eval.{name}].action_selection='{self.action_selection}' " - f"must be one of {valid_action_selections}" + f"[eval.{name}].action_selection='{self.action_selection}' must be one of {valid_action_selections}" ) # -- Config hooks --------------------------------------------------- @@ -172,9 +171,11 @@ def _run_rollout_loop(self, vecenv, policy, args) -> dict: with torch.no_grad(): ob_t = torch.as_tensor(obs).to(device) logits, _ = policy.forward_eval(ob_t, state) - action, _, _, cont_action = pufferlib.pytorch.sample_logits(logits, action_selection=self.action_selection, env_continuous=env_continuous, policy=policy) + action, _, _, cont_action = pufferlib.pytorch.sample_logits( + logits, action_selection=self.action_selection, env_continuous=env_continuous, policy=policy + ) - if env_continuous and not policy.is_continuous: # TODO check with trajectory + if env_continuous and not policy.is_continuous: # TODO check with trajectory cont_action = cont_action.cpu().numpy().reshape(vecenv.action_space.shape) obs, _, terminals, truncations, infos = vecenv.step(cont_action) else: @@ -540,9 +541,11 @@ def _render_pass_html(self, vecenv, policy, args) -> list: with torch.no_grad(): ob_t = torch.as_tensor(ob).to(device) logits, _ = policy.forward_eval(ob_t, state) - action, _, _, cont_action = pufferlib.pytorch.sample_logits(logits, action_selection=self.action_selection, env_continuous=env_continuous, policy=policy) + action, _, _, cont_action = pufferlib.pytorch.sample_logits( + logits, action_selection=self.action_selection, env_continuous=env_continuous, policy=policy + ) - if env_continuous and not policy.is_continuous: # TODO check with trajectory + if env_continuous and not policy.is_continuous: # TODO check with trajectory cont_action = cont_action.cpu().numpy().reshape(vecenv.action_space.shape) ob, _, terminals, truncations, infos = vec.step(cont_action) else: @@ -705,7 +708,9 @@ def _render_pass_obs(self, vecenv, policy, args) -> list: ob_t = torch.as_tensor(ob).to(device) logits, value = policy.forward_eval(ob_t, state) pool_outputs = pool_method(ob_t, state) if pool_method is not None else {} - action, logprob, entropy, cont_action = pufferlib.pytorch.sample_logits(logits, action_selection=self.action_selection, env_continuous=env_continuous, policy=policy) + action, logprob, entropy, cont_action = pufferlib.pytorch.sample_logits( + logits, action_selection=self.action_selection, env_continuous=env_continuous, policy=policy + ) pool_outputs = {k: v.cpu().numpy().astype(np.int16, copy=False) for k, v in pool_outputs.items()} if pool_hist is None and pool_outputs: pool_hist = { @@ -716,7 +721,7 @@ def _render_pass_obs(self, vecenv, policy, args) -> list: for k, values in pool_outputs.items() } - if env_continuous and not policy.is_continuous: # TODO check with trajectory + if env_continuous and not policy.is_continuous: # TODO check with trajectory raw_action = cont_action.cpu().numpy().reshape(vec.action_space.shape) clipped_action = raw_action ob, _, _, _, step_infos = vec.step(clipped_action) @@ -726,7 +731,7 @@ def _render_pass_obs(self, vecenv, policy, args) -> list: if isinstance(logits, torch.distributions.Normal): clipped_action = np.clip(raw_action, vec.action_space.low, vec.action_space.high) ob, _, _, _, step_infos = vec.step(clipped_action) - + if isinstance(logits, torch.distributions.Normal): policy_outputs = { "mean": logits.loc.cpu().numpy().reshape(vec.action_space.shape), @@ -785,7 +790,7 @@ def _render_pass_obs(self, vecenv, policy, args) -> list: np.asarray(policy_outputs[start_obs_index:end_obs_index], dtype=np.float32).copy() ) start_obs_index = end_obs_index - + for d in self._flatten_infos(step_infos): if isinstance(d, dict) and d.get("summary_type") == "completed_episode": batch_summary = d @@ -894,9 +899,11 @@ def _render_view(self, vecenv, target_env, policy, args, view_idx, out_dir, view with torch.no_grad(): ob_t = torch.as_tensor(ob).to(device) logits, _ = policy.forward_eval(ob_t, state) - action, _, _, cont_action = pufferlib.pytorch.sample_logits(logits, action_selection=self.action_selection, env_continuous=env_continuous, policy=policy) + action, _, _, cont_action = pufferlib.pytorch.sample_logits( + logits, action_selection=self.action_selection, env_continuous=env_continuous, policy=policy + ) - if env_continuous and not policy.is_continuous: # TODO check with trajectory + if env_continuous and not policy.is_continuous: # TODO check with trajectory cont_action = cont_action.cpu().numpy().reshape(vecenv.action_space.shape) ob, _, _, _, _ = vecenv.step(cont_action) else: diff --git a/pufferlib/ocean/torch.py b/pufferlib/ocean/torch.py index 02d9ab89bb..2db50d37c5 100644 --- a/pufferlib/ocean/torch.py +++ b/pufferlib/ocean/torch.py @@ -333,12 +333,32 @@ def __init__( action_type: str, ): super().__init__() - - self.register_buffer("JERK_LONG", torch.tensor((-15.0, -4.0, 0.0, 4.0), dtype=torch.float32, requires_grad=False), persistent=False) - self.register_buffer("JERK_LAT", torch.tensor((-4.0, 0.0, 4.0), dtype=torch.float32, requires_grad=False), persistent=False) - self.register_buffer("ACCELERATION_VALUES", torch.tensor((-4.0000, -2.6670, -1.3330, -0.0000, 1.3330, 2.6670, 4.0000), dtype=torch.float32, requires_grad=False), persistent=False) - self.register_buffer("STEERING_VALUES", torch.tensor((-0.667, -0.500, -0.333, -0.167, 0.000, 0.167, 0.333, 0.500, 0.667), dtype=torch.float32, requires_grad=False), persistent=False) + self.register_buffer( + "JERK_LONG", + torch.tensor((-15.0, -4.0, 0.0, 4.0), dtype=torch.float32, requires_grad=False), + persistent=False, + ) + self.register_buffer( + "JERK_LAT", torch.tensor((-4.0, 0.0, 4.0), dtype=torch.float32, requires_grad=False), persistent=False + ) + + self.register_buffer( + "ACCELERATION_VALUES", + torch.tensor( + (-4.0000, -2.6670, -1.3330, -0.0000, 1.3330, 2.6670, 4.0000), dtype=torch.float32, requires_grad=False + ), + persistent=False, + ) + self.register_buffer( + "STEERING_VALUES", + torch.tensor( + (-0.667, -0.500, -0.333, -0.167, 0.000, 0.167, 0.333, 0.500, 0.667), + dtype=torch.float32, + requires_grad=False, + ), + persistent=False, + ) if env.dynamics_model == "jerk": action_long, action_lat = self.JERK_LONG, self.JERK_LAT @@ -402,7 +422,9 @@ def __init__( self.critic_backbone = DriveBackbone(**backbone_args) # Setup action and value heads - self.is_continuous = action_type == "continuous" # TODO Check if what the `"trajectory"`, `"trajectory_frenet"`, `"trajectory_jerk" features do and if they are considered continuous. + self.is_continuous = ( + action_type == "continuous" + ) # TODO Check if what the `"trajectory"`, `"trajectory_frenet"`, `"trajectory_jerk" features do and if they are considered continuous. if self.is_continuous: self.atn_dim = (env.single_action_space.shape[0],) * 2 else: @@ -489,7 +511,7 @@ def decode_actions(self, hidden): value = self.critic_head(hidden) return action, value - + def discrete_actions_to_continuous(self, actions): return self.action_table[actions.long()] diff --git a/pufferlib/pufferl.py b/pufferlib/pufferl.py index c8456d9c12..833fc6f278 100644 --- a/pufferlib/pufferl.py +++ b/pufferlib/pufferl.py @@ -404,7 +404,9 @@ def evaluate(self): logits, value = self.policy.forward_eval(o_device, state) logits = logits_to_float(logits) - action, logprob, _, cont_action = pufferlib.pytorch.sample_logits(logits, env_continuous=self.env_continuous, policy=self.policy) + action, logprob, _, cont_action = pufferlib.pytorch.sample_logits( + logits, env_continuous=self.env_continuous, policy=self.policy + ) if config["normalize_rewards"]: r = torch.sign(r) * torch.log1p(torch.abs(r)) @@ -446,7 +448,6 @@ def evaluate(self): self.ep_lengths[env_id] = 0 self.free_idx += num_full self.full_rows += num_full - profile("eval_misc", epoch) for i in info: @@ -460,7 +461,7 @@ def evaluate(self): profile("env", epoch) - if self.env_continuous and not self.policy.is_continuous: # TODO check with trajectory + if self.env_continuous and not self.policy.is_continuous: # TODO check with trajectory cont_action = cont_action.cpu().numpy() self.vecenv.send(cont_action.squeeze(0)) else: @@ -1854,12 +1855,14 @@ def mine_failures(env_name, args=None): o_t = torch.as_tensor(obs_arr).to(device) state = {"reward": None, "done": None, "env_id": None, "mask": None} logits, _ = policy.forward_eval(o_t, state) - action, _, _, cont_action = pufferlib.pytorch.sample_logits(logits, env_continuous=env_continuous, policy=policy) + action, _, _, cont_action = pufferlib.pytorch.sample_logits( + logits, env_continuous=env_continuous, policy=policy + ) action = action.cpu().numpy() if action.ndim == 1 and len(vecenv.single_action_space.shape) >= 1: action = action.reshape(-1, *vecenv.single_action_space.shape) - if env_continuous and not policy.is_continuous: # TODO check with trajectory + if env_continuous and not policy.is_continuous: # TODO check with trajectory cont_action = cont_action.cpu().numpy() vecenv.send(cont_action.squeeze(0)) else: diff --git a/pufferlib/pytorch.py b/pufferlib/pytorch.py index 871d8a192e..f567a8a8c8 100644 --- a/pufferlib/pytorch.py +++ b/pufferlib/pytorch.py @@ -21,8 +21,8 @@ } ACTION_SELECT_SAMPLE = "sample" -ACTION_SELECT_MODE = "mode" # argmax -ACTION_SELECT_MEAN = "mean" # probability-weighted continuous mean +ACTION_SELECT_MODE = "mode" # argmax +ACTION_SELECT_MEAN = "mean" # probability-weighted continuous mean LITTLE_BYTE_ORDER = sys.byteorder == "little" @@ -190,7 +190,9 @@ def entropy_probs(logits, probs): return -p_log_p.sum(-1) -def sample_logits(logits, action=None, action_selection=ACTION_SELECT_SAMPLE, env_continuous=None, policy=None): # TODO discrete continuous +def sample_logits( + logits, action=None, action_selection=ACTION_SELECT_SAMPLE, env_continuous=None, policy=None +): # TODO discrete continuous is_discrete = isinstance(logits, torch.Tensor) if action_selection == ACTION_SELECT_MEAN and not (env_continuous and not policy.is_continuous): raise ValueError("action_selection='mean' requires a discrete policy on a continuous env") @@ -234,7 +236,7 @@ def sample_logits(logits, action=None, action_selection=ACTION_SELECT_SAMPLE, en if env_continuous and not policy.is_continuous: if action_selection == ACTION_SELECT_MEAN: cont_actions = policy.discrete_probs_to_continuous_mean(probs) - else: # Mode and sample we already have an action from the code above. + else: # Mode and sample we already have an action from the code above. cont_actions = policy.discrete_actions_to_continuous(action) if is_discrete: return action.squeeze(0), logprob.squeeze(0), logits_entropy.squeeze(0), cont_actions.squeeze(0) diff --git a/scripts/0_build_puffer.sh b/scripts/0_build_puffer.sh index dd1bbaeae0..4f68a5d85b 100644 --- a/scripts/0_build_puffer.sh +++ b/scripts/0_build_puffer.sh @@ -20,4 +20,4 @@ python setup.py build_ext --inplace --force end=`date +%s` runtime=$((end-start)) echo "END TIME: $(date)" -echo "Runtime: ${runtime}" \ No newline at end of file +echo "Runtime: ${runtime}" diff --git a/scripts/0_eval_model_slurm.sh b/scripts/0_eval_model_slurm.sh index c8d49b7692..6a93d91a99 100644 --- a/scripts/0_eval_model_slurm.sh +++ b/scripts/0_eval_model_slurm.sh @@ -20,4 +20,4 @@ source .venv/bin/activate end=`date +%s` runtime=$((end-start)) echo "END TIME: $(date)" -echo "Runtime: ${runtime}" \ No newline at end of file +echo "Runtime: ${runtime}" diff --git a/scripts/0_train_model_slurm.sh b/scripts/0_train_model_slurm.sh index c59dfa2688..7eadf39562 100644 --- a/scripts/0_train_model_slurm.sh +++ b/scripts/0_train_model_slurm.sh @@ -76,4 +76,4 @@ torchrun --standalone --nnodes=1 --nproc-per-node=8 --max_restarts=0 --start-met end=$(date +%s) runtime=$((end-start)) echo "END TIME: $(date)" -echo "Runtime: ${runtime}" \ No newline at end of file +echo "Runtime: ${runtime}" diff --git a/scripts/cluster_configs/bernhard_1.yaml b/scripts/cluster_configs/bernhard_1.yaml index 90b8798bd2..0b563b62c2 100644 --- a/scripts/cluster_configs/bernhard_1.yaml +++ b/scripts/cluster_configs/bernhard_1.yaml @@ -4,4 +4,4 @@ nodes: 1 gpus: 1 cpus: 28 mem: 128gb -time: 2880 # minutes \ No newline at end of file +time: 2880 # minutes diff --git a/scripts/cluster_configs/bernhard_8.yaml b/scripts/cluster_configs/bernhard_8.yaml index 5bd0b7f4dc..7b38582aa6 100644 --- a/scripts/cluster_configs/bernhard_8.yaml +++ b/scripts/cluster_configs/bernhard_8.yaml @@ -4,4 +4,4 @@ nodes: 1 gpus: 8 cpus: 224 mem: 1024gb -time: 5760 # minutes \ No newline at end of file +time: 5760 # minutes From a0bad3f8c80eea43898ed7ea617fb8adc7579b11 Mon Sep 17 00:00:00 2001 From: Bernhard Jaeger Date: Thu, 23 Jul 2026 21:07:22 +0200 Subject: [PATCH 26/35] fix bug where policy would crash due to ddp naming --- pufferlib/pufferl.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/pufferlib/pufferl.py b/pufferlib/pufferl.py index 833fc6f278..0624311a97 100644 --- a/pufferlib/pufferl.py +++ b/pufferlib/pufferl.py @@ -129,8 +129,10 @@ def __init__(self, config, vecenv, policy, logger=None): self.env_continuous = isinstance(vecenv.single_action_space, pufferlib.spaces.Box) obs_space = vecenv.single_observation_space - if self.env_continuous and not policy.is_continuous: - action_shape = (len(policy.atn_dim),) + # Custom policy attributes live on the base module, not the DDP/compile wrapper. + unwrapped_policy = base_policy(policy) + if self.env_continuous and not unwrapped_policy.is_continuous: + action_shape = (len(unwrapped_policy.atn_dim),) action_dtype = torch.int32 else: action_shape = vecenv.single_action_space @@ -405,7 +407,7 @@ def evaluate(self): logits, value = self.policy.forward_eval(o_device, state) logits = logits_to_float(logits) action, logprob, _, cont_action = pufferlib.pytorch.sample_logits( - logits, env_continuous=self.env_continuous, policy=self.policy + logits, env_continuous=self.env_continuous, policy=self.uncompiled_policy ) if config["normalize_rewards"]: r = torch.sign(r) * torch.log1p(torch.abs(r)) @@ -461,7 +463,7 @@ def evaluate(self): profile("env", epoch) - if self.env_continuous and not self.policy.is_continuous: # TODO check with trajectory + if self.env_continuous and not self.uncompiled_policy.is_continuous: # TODO check with trajectory cont_action = cont_action.cpu().numpy() self.vecenv.send(cont_action.squeeze(0)) else: From 55ef51006a040723045be7f20236454ae976f45e Mon Sep 17 00:00:00 2001 From: Bernhard Jaeger Date: Thu, 23 Jul 2026 21:21:24 +0200 Subject: [PATCH 27/35] Fix for smoke test --- pufferlib/pufferl.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pufferlib/pufferl.py b/pufferlib/pufferl.py index 0624311a97..6e32e56482 100644 --- a/pufferlib/pufferl.py +++ b/pufferlib/pufferl.py @@ -135,7 +135,7 @@ def __init__(self, config, vecenv, policy, logger=None): action_shape = (len(unwrapped_policy.atn_dim),) action_dtype = torch.int32 else: - action_shape = vecenv.single_action_space + action_shape = vecenv.single_action_space.shape action_dtype = pufferlib.pytorch.numpy_to_torch_dtype_dict[vecenv.single_action_space.dtype] total_agents = vecenv.num_agents self.total_agents = total_agents From ba1aefb934b5b5649136ee8cac59f49c98bd3386 Mon Sep 17 00:00:00 2001 From: Bernhard Jaeger Date: Fri, 24 Jul 2026 18:25:38 +0200 Subject: [PATCH 28/35] Jobs for 8 GPUs and an eval array. --- scripts/1_train_model_slurm.sh | 14 ++++++++++---- scripts/2_train_model_multinode_slurm.sh | 14 ++++++++------ 2 files changed, 18 insertions(+), 10 deletions(-) diff --git a/scripts/1_train_model_slurm.sh b/scripts/1_train_model_slurm.sh index 4c9fe9f8bc..75c137d6e7 100644 --- a/scripts/1_train_model_slurm.sh +++ b/scripts/1_train_model_slurm.sh @@ -6,17 +6,22 @@ #SBATCH --gres gpu:8 #SBATCH --mem=1007G #SBATCH --cpus-per-task 144 -#SBATCH --output /home/bjaeger/PufferDrive/experiments/k_nightly_0006/log_%a_%A.out -#SBATCH --error /home/bjaeger/PufferDrive/experiments/k_nightly_0006/log_%a_%A.err +#SBATCH --output /home/bjaeger/PufferDrive/experiments/k_nightly_0007/log_%a_%A.out +#SBATCH --error /home/bjaeger/PufferDrive/experiments/k_nightly_0007/log_%a_%A.err #SBATCH --partition dev +#SBATCH --array=0-3 # print info about current job echo "START TIME: $(date)" start=$(date +%s) -export RUN_NAME=k_nightly_0006 +export RUN_NAME=k_nightly_0007 echo ${RUN_NAME} +# Seed each array task deterministically: 1000 * array task id +SEED=$((1000 * SLURM_ARRAY_TASK_ID)) +echo "SLURM_ARRAY_TASK_ID=${SLURM_ARRAY_TASK_ID} -> train.seed=${SEED}" + # TODO could try to tune these. 1 Is probably best since Puffer parallelizes across all cores. export NUMEXPR_NUM_THREADS=1 export MKL_NUM_THREADS=1 @@ -32,7 +37,7 @@ torchrun --standalone --nnodes=1 --nproc-per-node=8 --max_restarts=0 --start-met wandb_group=emerge_ \ train.data_dir=/home/bjaeger/PufferDrive/experiments/${RUN_NAME} \ env.map_dir=/home/bjaeger/PufferDrive/pufferlib/resources/drive/binaries/carla \ - train.name=${RUN_NAME} \ + train.name=${RUN_NAME}_${SEED} \ train.total_timesteps=100000000000 \ vec.num_envs=16 \ +eval.map_dir=/home/bjaeger/data/nuPlan/PufferDrive \ @@ -41,6 +46,7 @@ torchrun --standalone --nnodes=1 --nproc-per-node=8 --max_restarts=0 --start-met train.minibatch_size=131072 \ train.precision=bfloat16 \ eval.validation_gigaflow.render_backend=obs_html \ + train.seed=${SEED} \ tb=True diff --git a/scripts/2_train_model_multinode_slurm.sh b/scripts/2_train_model_multinode_slurm.sh index a453ffe7d3..445a6938b0 100644 --- a/scripts/2_train_model_multinode_slurm.sh +++ b/scripts/2_train_model_multinode_slurm.sh @@ -1,13 +1,13 @@ #!/bin/bash #SBATCH --job-name train_puffer -#SBATCH --nodes 2 # Number of nodes requested +#SBATCH --nodes 8 # Number of nodes requested #SBATCH --ntasks-per-node 1 # Run 1 srun task per node (which fires up torchrun) #SBATCH --gres gpu:8 # GPUs per node #SBATCH --cpus-per-task 144 #SBATCH --mem=1007G -#SBATCH --time 2-00:00 -#SBATCH --output /home/bjaeger/nightly_PufferDrive/experiments/k_nightly_0006/log_%a_%A.out -#SBATCH --error /home/bjaeger/nightly_PufferDrive/experiments/k_nightly_0006/log_%a_%A.err +#SBATCH --time 3-00:00 +#SBATCH --output /home/bjaeger/nightly_PufferDrive/experiments/k_scaled_0000/log_%a_%A.out +#SBATCH --error /home/bjaeger/nightly_PufferDrive/experiments/k_scaled_0000/log_%a_%A.err #SBATCH --partition dev # Set up PyTorch Distributed Rendezvous parameters from Slurm variables @@ -20,7 +20,7 @@ echo "Master node: ${MASTER_ADDR}:${MASTER_PORT}" echo "Total nodes: ${WORLD_SIZE}" start=$(date +%s) -export RUN_NAME=k_nightly_0006 +export RUN_NAME=k_scaled_0000 echo ${RUN_NAME} # Thread limit limits CPU thrashing across worker environments @@ -48,13 +48,15 @@ srun torchrun \ train.data_dir=/home/bjaeger/PufferDrive/experiments/${RUN_NAME} \ env.map_dir=/home/bjaeger/PufferDrive/pufferlib/resources/drive/binaries/carla \ train.name=${RUN_NAME} \ - train.total_timesteps=100000000000 \ + train.total_timesteps=1000000000000 \ vec.num_envs=16 \ +eval.map_dir=/home/bjaeger/data/nuPlan/PufferDrive \ train.compile=True \ train.max_minibatch_size=131072 \ train.minibatch_size=131072 \ train.precision=bfloat16 \ + eval.validation_gigaflow.render_backend=obs_html \ + train.seed=0 \ tb=True end=$(date +%s) From 9c1adcb9880b70bc640cb3d6464cf86903116b28 Mon Sep 17 00:00:00 2001 From: Bernhard Jaeger Date: Fri, 24 Jul 2026 18:29:26 +0200 Subject: [PATCH 29/35] MInor changes to script. --- scripts/1_train_model_slurm.sh | 2 +- scripts/2_train_model_multinode_slurm.sh | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/scripts/1_train_model_slurm.sh b/scripts/1_train_model_slurm.sh index 75c137d6e7..077cb66a77 100644 --- a/scripts/1_train_model_slurm.sh +++ b/scripts/1_train_model_slurm.sh @@ -9,7 +9,7 @@ #SBATCH --output /home/bjaeger/PufferDrive/experiments/k_nightly_0007/log_%a_%A.out #SBATCH --error /home/bjaeger/PufferDrive/experiments/k_nightly_0007/log_%a_%A.err #SBATCH --partition dev -#SBATCH --array=0-3 +#SBATCH --array=0-2 # print info about current job echo "START TIME: $(date)" diff --git a/scripts/2_train_model_multinode_slurm.sh b/scripts/2_train_model_multinode_slurm.sh index 445a6938b0..bb5cec2c32 100644 --- a/scripts/2_train_model_multinode_slurm.sh +++ b/scripts/2_train_model_multinode_slurm.sh @@ -6,8 +6,8 @@ #SBATCH --cpus-per-task 144 #SBATCH --mem=1007G #SBATCH --time 3-00:00 -#SBATCH --output /home/bjaeger/nightly_PufferDrive/experiments/k_scaled_0000/log_%a_%A.out -#SBATCH --error /home/bjaeger/nightly_PufferDrive/experiments/k_scaled_0000/log_%a_%A.err +#SBATCH --output /home/bjaeger/PufferDrive/experiments/k_scaled_0000/log_%a_%A.out +#SBATCH --error /home/bjaeger/PufferDrive/experiments/k_scaled_0000/log_%a_%A.err #SBATCH --partition dev # Set up PyTorch Distributed Rendezvous parameters from Slurm variables From 36c06e355eec382189947c052a6309c34239e764 Mon Sep 17 00:00:00 2001 From: Bernhard Jaeger Date: Mon, 27 Jul 2026 12:58:21 +0200 Subject: [PATCH 30/35] Test 64 GPU with 100B. --- scripts/2_train_model_multinode_slurm.sh | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/scripts/2_train_model_multinode_slurm.sh b/scripts/2_train_model_multinode_slurm.sh index bb5cec2c32..0088ffaded 100644 --- a/scripts/2_train_model_multinode_slurm.sh +++ b/scripts/2_train_model_multinode_slurm.sh @@ -6,8 +6,8 @@ #SBATCH --cpus-per-task 144 #SBATCH --mem=1007G #SBATCH --time 3-00:00 -#SBATCH --output /home/bjaeger/PufferDrive/experiments/k_scaled_0000/log_%a_%A.out -#SBATCH --error /home/bjaeger/PufferDrive/experiments/k_scaled_0000/log_%a_%A.err +#SBATCH --output /home/bjaeger/PufferDrive/experiments/k_scaled_0001/log_%a_%A.out +#SBATCH --error /home/bjaeger/PufferDrive/experiments/k_scaled_0001/log_%a_%A.err #SBATCH --partition dev # Set up PyTorch Distributed Rendezvous parameters from Slurm variables @@ -20,7 +20,7 @@ echo "Master node: ${MASTER_ADDR}:${MASTER_PORT}" echo "Total nodes: ${WORLD_SIZE}" start=$(date +%s) -export RUN_NAME=k_scaled_0000 +export RUN_NAME=k_scaled_0001 echo ${RUN_NAME} # Thread limit limits CPU thrashing across worker environments @@ -48,7 +48,7 @@ srun torchrun \ train.data_dir=/home/bjaeger/PufferDrive/experiments/${RUN_NAME} \ env.map_dir=/home/bjaeger/PufferDrive/pufferlib/resources/drive/binaries/carla \ train.name=${RUN_NAME} \ - train.total_timesteps=1000000000000 \ + train.total_timesteps=100000000000 \ vec.num_envs=16 \ +eval.map_dir=/home/bjaeger/data/nuPlan/PufferDrive \ train.compile=True \ From 3ec5a0b308b613d357ba0cf3d492ae399a6cdae0 Mon Sep 17 00:00:00 2001 From: Bernhard Jaeger Date: Mon, 27 Jul 2026 13:06:01 +0200 Subject: [PATCH 31/35] k_scaled_0002, 100B, learning rate * 8 --- scripts/2_train_model_multinode_slurm.sh | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/scripts/2_train_model_multinode_slurm.sh b/scripts/2_train_model_multinode_slurm.sh index 0088ffaded..25e193631a 100644 --- a/scripts/2_train_model_multinode_slurm.sh +++ b/scripts/2_train_model_multinode_slurm.sh @@ -6,8 +6,8 @@ #SBATCH --cpus-per-task 144 #SBATCH --mem=1007G #SBATCH --time 3-00:00 -#SBATCH --output /home/bjaeger/PufferDrive/experiments/k_scaled_0001/log_%a_%A.out -#SBATCH --error /home/bjaeger/PufferDrive/experiments/k_scaled_0001/log_%a_%A.err +#SBATCH --output /home/bjaeger/PufferDrive/experiments/k_scaled_0002/log_%a_%A.out +#SBATCH --error /home/bjaeger/PufferDrive/experiments/k_scaled_0002/log_%a_%A.err #SBATCH --partition dev # Set up PyTorch Distributed Rendezvous parameters from Slurm variables @@ -20,7 +20,7 @@ echo "Master node: ${MASTER_ADDR}:${MASTER_PORT}" echo "Total nodes: ${WORLD_SIZE}" start=$(date +%s) -export RUN_NAME=k_scaled_0001 +export RUN_NAME=k_scaled_0002 echo ${RUN_NAME} # Thread limit limits CPU thrashing across worker environments @@ -49,6 +49,7 @@ srun torchrun \ env.map_dir=/home/bjaeger/PufferDrive/pufferlib/resources/drive/binaries/carla \ train.name=${RUN_NAME} \ train.total_timesteps=100000000000 \ + train.learning_rate=0.004 \ vec.num_envs=16 \ +eval.map_dir=/home/bjaeger/data/nuPlan/PufferDrive \ train.compile=True \ From ab33574e15f92030f6b30a8e14abd9379f826c35 Mon Sep 17 00:00:00 2001 From: Bernhard Jaeger Date: Mon, 27 Jul 2026 14:06:06 +0200 Subject: [PATCH 32/35] Fix unit tests. --- pufferlib/ocean/benchmark/evaluators/base.py | 6 ++-- tests/unit_tests/test_eval_manager.py | 33 +++++++++++++++++--- 2 files changed, 33 insertions(+), 6 deletions(-) diff --git a/pufferlib/ocean/benchmark/evaluators/base.py b/pufferlib/ocean/benchmark/evaluators/base.py index 61c2ed9bc2..e87a6868e0 100644 --- a/pufferlib/ocean/benchmark/evaluators/base.py +++ b/pufferlib/ocean/benchmark/evaluators/base.py @@ -872,8 +872,6 @@ def _render_view(self, vecenv, target_env, policy, args, view_idx, out_dir, view num_scenarios = int(eval_cfg["render_num_scenarios"]) max_steps = int(eval_cfg["render_max_steps"]) - saved_cwd = os.getcwd() - os.chdir(out_dir) # Glob by full view_suffix (= step_suffix + view marker) so we get # only this-view's mp4s — not files written by a prior view in the # same render pass, which would otherwise duplicate in all_paths. @@ -881,7 +879,11 @@ def _render_view(self, vecenv, target_env, policy, args, view_idx, out_dir, view # matches bev files but not the bare sim_state ones. view_glob = f"*{view_suffix}.mp4" env_continuous = isinstance(vecenv.single_action_space, pufferlib.spaces.Box) + # chdir must be inside the try so the finally always restores the cwd, + # even if a later statement raises before the loop starts. + saved_cwd = os.getcwd() try: + os.chdir(out_dir) state = self._init_lstm_state(num_agents, policy, device, args) scenarios_processed = 0 while scenarios_processed < num_scenarios: diff --git a/tests/unit_tests/test_eval_manager.py b/tests/unit_tests/test_eval_manager.py index 07628fb8bf..c7116d341b 100644 --- a/tests/unit_tests/test_eval_manager.py +++ b/tests/unit_tests/test_eval_manager.py @@ -470,10 +470,12 @@ def test_rollout_zeros_lstm_state_per_agent_on_done(monkeypatch): """Per-agent LSTM reset on terminations or truncations. Either signal means 'episode over, env reset' — the agent's next obs is from a fresh scenario and stale recurrent memory would bias the policy.""" + import gymnasium import numpy as np import torch import pufferlib.pytorch + import pufferlib.spaces from pufferlib.ocean.benchmark.evaluators.base import Evaluator state = {"lstm_h": torch.ones(4, 8), "lstm_c": torch.ones(4, 8)} @@ -492,7 +494,10 @@ def _should_stop(self, args, infos_collected, steps): class _Vec: observation_space = type("S", (), {"shape": (4, 6)})() - action_space = type("A", (), {"shape": (4,), "low": -1.0, "high": 1.0})() + # Discrete env, like drive "jerk": per-agent MultiDiscrete([12]) (not a + # Box -> env_continuous is False), joint action_space is Box(shape=(4, 1)). + single_action_space = gymnasium.spaces.MultiDiscrete([12]) + action_space = pufferlib.spaces.joint_space(single_action_space, 4) reset_obs = np.zeros((4, 6), dtype=np.float32) def step(self, action): @@ -500,6 +505,8 @@ def step(self, action): return self.reset_obs, np.zeros(4), np.array([0, 1, 0, 0]), np.array([1, 0, 1, 0]), [] class _Policy: + is_continuous = False + def forward_eval(self, ob, state): return torch.zeros(ob.shape[0], 1), None @@ -508,7 +515,12 @@ def forward_eval(self, ob, state): monkeypatch.setattr( pufferlib.pytorch, "sample_logits", - lambda logits, deterministic=True: (torch.zeros(4, dtype=torch.long), None, None), + lambda logits, action=None, action_selection=pufferlib.pytorch.ACTION_SELECT_MODE, env_continuous=None, policy=None: ( + torch.zeros(4, dtype=torch.long), + None, + None, + None, + ), ) args = {"train": {"device": "cpu", "use_rnn": True}, "env": {}} @@ -715,6 +727,9 @@ def test_render_pass_per_evaluator_subdir_and_step_glob(tmp_path, monkeypatch): Mocks the env construction + render so we can control the filenames the fake env produces. """ + import gymnasium + + import pufferlib.spaces import pufferlib.vector from pufferlib.ocean.benchmark.evaluators.base import Evaluator @@ -740,7 +755,10 @@ def close_client(self, env_idx): class _FakeVec: observation_space = type("S", (), {"shape": (2, 4)})() - action_space = type("A", (), {"shape": (2,), "low": -1.0, "high": 1.0})() + # Discrete env (2 agents): per-agent MultiDiscrete([12]) (not a Box -> + # env_continuous is False), joint action_space is Box(shape=(2, 1)). + single_action_space = gymnasium.spaces.MultiDiscrete([12]) + action_space = pufferlib.spaces.joint_space(single_action_space, 2) def reset(self): import numpy as np @@ -781,13 +799,20 @@ def close(self): import pufferlib.pytorch class _Policy: + is_continuous = False + def forward_eval(self, ob, state): return torch.zeros(ob.shape[0], 1), None monkeypatch.setattr( pufferlib.pytorch, "sample_logits", - lambda logits, deterministic=True: (torch.zeros(logits.shape[0], dtype=torch.long), None, None), + lambda logits, action=None, action_selection=pufferlib.pytorch.ACTION_SELECT_MODE, env_continuous=None, policy=None: ( + torch.zeros(logits.shape[0], dtype=torch.long), + None, + None, + None, + ), ) class _Ev(Evaluator): From 37b59c29880a82dfc9bd085fbf46b55dc54443e8 Mon Sep 17 00:00:00 2001 From: Bernhard Jaeger Date: Mon, 27 Jul 2026 14:10:30 +0200 Subject: [PATCH 33/35] formatting. --- bernhard_notes.txt | 4 +++- tests/unit_tests/test_eval_manager.py | 12 ++++++++++-- 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/bernhard_notes.txt b/bernhard_notes.txt index 765a5e3205..ad2fa6474e 100644 --- a/bernhard_notes.txt +++ b/bernhard_notes.txt @@ -1,6 +1,8 @@ git merge origin/3.0 git push origin bernhard +Run for formatting. +pre-commit run --all-files # Not in valentins config but in my config. # non_sdc_controller: policy @@ -107,7 +109,7 @@ lauch.json: "name"="PufferDrive Train", "type"="python", "request"="launch", - "program"="${workspaceFolder}/.venv/bin/puffer", + "program"="${workspaceFolder}pre-commit run --all-files/.venv/bin/puffer", "args"=[ "train", "puffer_drive", diff --git a/tests/unit_tests/test_eval_manager.py b/tests/unit_tests/test_eval_manager.py index c7116d341b..24e37698e8 100644 --- a/tests/unit_tests/test_eval_manager.py +++ b/tests/unit_tests/test_eval_manager.py @@ -515,7 +515,11 @@ def forward_eval(self, ob, state): monkeypatch.setattr( pufferlib.pytorch, "sample_logits", - lambda logits, action=None, action_selection=pufferlib.pytorch.ACTION_SELECT_MODE, env_continuous=None, policy=None: ( + lambda logits, + action=None, + action_selection=pufferlib.pytorch.ACTION_SELECT_MODE, + env_continuous=None, + policy=None: ( torch.zeros(4, dtype=torch.long), None, None, @@ -807,7 +811,11 @@ def forward_eval(self, ob, state): monkeypatch.setattr( pufferlib.pytorch, "sample_logits", - lambda logits, action=None, action_selection=pufferlib.pytorch.ACTION_SELECT_MODE, env_continuous=None, policy=None: ( + lambda logits, + action=None, + action_selection=pufferlib.pytorch.ACTION_SELECT_MODE, + env_continuous=None, + policy=None: ( torch.zeros(logits.shape[0], dtype=torch.long), None, None, From bcff09daa8417c5669ba127bad7b9280ef99d903 Mon Sep 17 00:00:00 2001 From: Bernhard Jaeger Date: Mon, 27 Jul 2026 14:34:06 +0200 Subject: [PATCH 34/35] Added unit test to check the new feature. --- tests/unit_tests/test_action_selection.py | 188 ++++++++++++++++++++++ 1 file changed, 188 insertions(+) create mode 100644 tests/unit_tests/test_action_selection.py diff --git a/tests/unit_tests/test_action_selection.py b/tests/unit_tests/test_action_selection.py new file mode 100644 index 0000000000..0627857a34 --- /dev/null +++ b/tests/unit_tests/test_action_selection.py @@ -0,0 +1,188 @@ +"""Unit tests for evaluation-time action selection (sample / mode / mean). + +The feature lets a user pick how eval turns policy logits into actions on a +*continuous* environment, for either a discrete or a continuous policy: + + - sample: draw from the policy distribution (stochastic) + - mode: argmax / Gaussian mean (deterministic "best guess") + - mean: probability-weighted continuous mean (discrete policy only) + +The logic lives in `pufferlib.pytorch.sample_logits`; the mode is chosen by +`Evaluator.action_selection` (config -> validated -> passed to sample_logits). +Both are covered here without an env/C-sim/GPU. +""" + +import pytest +import torch + +import pufferlib.pytorch as P + +# 3 discrete classes -> fixed 2D continuous embedding, mirroring Drive's +# action_table[num_classes, 2] (discrete action index -> (long, lat)). +ACTION_TABLE = torch.tensor([[-1.0, -1.0], [0.0, 0.0], [1.0, 1.0]]) + + +class _DiscretePolicy: + """Minimal stand-in for a discrete Drive policy on a continuous env.""" + + is_continuous = False + + def discrete_actions_to_continuous(self, actions): + return ACTION_TABLE[actions.long()] + + def discrete_probs_to_continuous_mean(self, probs): + return probs @ ACTION_TABLE.to(probs.dtype) + + +class _ContinuousPolicy: + is_continuous = True + + +def _discrete_logits(rows): + # Drive's discrete policy returns a tuple of one [batch, num_classes] + # tensor (torch.split output); mirror that so we exercise the real path. + return (torch.tensor(rows),) + + +# -------------------------------------------------------------------------- +# Discrete policy on a continuous env +# -------------------------------------------------------------------------- +def test_mode_selects_argmax_and_is_deterministic(): + # Row 0 argmax = class 2, row 1 argmax = class 0. + logits = _discrete_logits([[0.0, 0.0, 10.0], [10.0, 0.0, 0.0]]) + policy = _DiscretePolicy() + + action1, _, _, cont1 = P.sample_logits( + logits, action_selection=P.ACTION_SELECT_MODE, env_continuous=True, policy=policy + ) + action2, _, _, cont2 = P.sample_logits( + logits, action_selection=P.ACTION_SELECT_MODE, env_continuous=True, policy=policy + ) + + assert action1.reshape(-1).tolist() == [2, 0] + # Continuous action is the exact table row for the argmax class. + expected_cont = torch.stack([ACTION_TABLE[2], ACTION_TABLE[0]]) + assert torch.allclose(cont1.reshape(-1, 2), expected_cont) + # Deterministic: repeated calls are identical. + assert torch.equal(action1, action2) + assert torch.equal(cont1, cont2) + + +def test_mean_is_probability_weighted_and_differs_from_mode(): + # Uniform distribution over the 3 classes. + logits = _discrete_logits([[0.0, 0.0, 0.0]]) + policy = _DiscretePolicy() + + action_mean, _, _, cont_mean = P.sample_logits( + logits, action_selection=P.ACTION_SELECT_MEAN, env_continuous=True, policy=policy + ) + action_mode, _, _, cont_mode = P.sample_logits( + logits, action_selection=P.ACTION_SELECT_MODE, env_continuous=True, policy=policy + ) + + # mean == probability-weighted continuous mean == average of table rows == [0, 0]. + assert torch.allclose(cont_mean.reshape(-1), torch.zeros(2), atol=1e-6) + # mode picks the argmax row ([-1, -1] for a uniform tie) -> distinct from mean. + assert not torch.allclose(cont_mean.reshape(-1), cont_mode.reshape(-1)) + # The nominal discrete action reported for mean is the argmax (for logging). + assert action_mean.reshape(-1).tolist() == action_mode.reshape(-1).tolist() + + +def test_sample_is_stochastic_and_stays_in_support(): + # Flat distribution so multinomial explores every class. + logits = _discrete_logits([[0.0, 0.0, 0.0]]) + policy = _DiscretePolicy() + + torch.manual_seed(0) + seen_classes = set() + for _ in range(50): + action, _, _, cont = P.sample_logits( + logits, action_selection=P.ACTION_SELECT_SAMPLE, env_continuous=True, policy=policy + ) + idx = int(action.reshape(-1).item()) + seen_classes.add(idx) + assert idx in (0, 1, 2) + # Continuous action must be the table row of the *sampled* class. + assert torch.allclose(cont.reshape(-1), ACTION_TABLE[idx]) + # Genuinely sampling, not collapsing to argmax. + assert len(seen_classes) > 1 + + +# -------------------------------------------------------------------------- +# Continuous policy on a continuous env +# -------------------------------------------------------------------------- +def test_continuous_policy_mode_returns_gaussian_mean(): + loc = torch.tensor([[0.5, -0.5]]) + dist = torch.distributions.Normal(loc, torch.ones(1, 2)) + + action, _, _, cont = P.sample_logits( + dist, action_selection=P.ACTION_SELECT_MODE, env_continuous=True, policy=_ContinuousPolicy() + ) + + # A continuous policy already emits a continuous action -> no discrete->cont conversion. + assert cont is None + # Mode of a Gaussian is its mean (loc), deterministically. + assert torch.allclose(action.reshape(1, 2), loc) + + +def test_continuous_policy_sample_differs_from_mode(): + loc = torch.tensor([[0.5, -0.5]]) + dist = torch.distributions.Normal(loc, torch.ones(1, 2)) + + torch.manual_seed(0) + action_sample, _, _, _ = P.sample_logits( + dist, action_selection=P.ACTION_SELECT_SAMPLE, env_continuous=True, policy=_ContinuousPolicy() + ) + + assert not torch.allclose(action_sample.reshape(1, 2), loc) + + +# -------------------------------------------------------------------------- +# Guards: `mean` is only valid for a discrete policy on a continuous env +# -------------------------------------------------------------------------- +def test_mean_rejects_continuous_policy(): + dist = torch.distributions.Normal(torch.zeros(1, 2), torch.ones(1, 2)) + with pytest.raises(ValueError): + P.sample_logits( + dist, action_selection=P.ACTION_SELECT_MEAN, env_continuous=True, policy=_ContinuousPolicy() + ) + + +def test_mean_rejects_discrete_env(): + logits = _discrete_logits([[0.0, 0.0, 0.0]]) + with pytest.raises(ValueError): + P.sample_logits( + logits, action_selection=P.ACTION_SELECT_MEAN, env_continuous=False, policy=_DiscretePolicy() + ) + + +# -------------------------------------------------------------------------- +# Evaluator config plumbing: config -> validated self.action_selection +# -------------------------------------------------------------------------- +def _make_evaluator(action_selection=None): + from pufferlib.ocean.benchmark.evaluators.base import Evaluator + + class _Ev(Evaluator): + type_name = "_action_selection_test" + + def _should_stop(self, args, infos_collected, steps): + return True + + config = {} if action_selection is None else {"action_selection": action_selection} + return _Ev("atn_test", config, train_config={}) + + +def test_evaluator_defaults_to_mode(): + assert _make_evaluator().action_selection == P.ACTION_SELECT_MODE + + +@pytest.mark.parametrize( + "mode", [P.ACTION_SELECT_SAMPLE, P.ACTION_SELECT_MODE, P.ACTION_SELECT_MEAN] +) +def test_evaluator_accepts_valid_action_selection(mode): + assert _make_evaluator(mode).action_selection == mode + + +def test_evaluator_rejects_invalid_action_selection(): + with pytest.raises(ValueError): + _make_evaluator("banana") From 2809f768db01323a1a21c5f924e1dee92b3a1b6b Mon Sep 17 00:00:00 2001 From: Bernhard Jaeger Date: Mon, 27 Jul 2026 14:36:05 +0200 Subject: [PATCH 35/35] Formatting --- tests/unit_tests/test_action_selection.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/tests/unit_tests/test_action_selection.py b/tests/unit_tests/test_action_selection.py index 0627857a34..89afea37e1 100644 --- a/tests/unit_tests/test_action_selection.py +++ b/tests/unit_tests/test_action_selection.py @@ -143,17 +143,13 @@ def test_continuous_policy_sample_differs_from_mode(): def test_mean_rejects_continuous_policy(): dist = torch.distributions.Normal(torch.zeros(1, 2), torch.ones(1, 2)) with pytest.raises(ValueError): - P.sample_logits( - dist, action_selection=P.ACTION_SELECT_MEAN, env_continuous=True, policy=_ContinuousPolicy() - ) + P.sample_logits(dist, action_selection=P.ACTION_SELECT_MEAN, env_continuous=True, policy=_ContinuousPolicy()) def test_mean_rejects_discrete_env(): logits = _discrete_logits([[0.0, 0.0, 0.0]]) with pytest.raises(ValueError): - P.sample_logits( - logits, action_selection=P.ACTION_SELECT_MEAN, env_continuous=False, policy=_DiscretePolicy() - ) + P.sample_logits(logits, action_selection=P.ACTION_SELECT_MEAN, env_continuous=False, policy=_DiscretePolicy()) # -------------------------------------------------------------------------- @@ -176,9 +172,7 @@ def test_evaluator_defaults_to_mode(): assert _make_evaluator().action_selection == P.ACTION_SELECT_MODE -@pytest.mark.parametrize( - "mode", [P.ACTION_SELECT_SAMPLE, P.ACTION_SELECT_MODE, P.ACTION_SELECT_MEAN] -) +@pytest.mark.parametrize("mode", [P.ACTION_SELECT_SAMPLE, P.ACTION_SELECT_MODE, P.ACTION_SELECT_MEAN]) def test_evaluator_accepts_valid_action_selection(mode): assert _make_evaluator(mode).action_selection == mode