Add come updates for Neurips paper (#4)

* scenarionet training * wandb * train utils * fix callback * run PPO * use pg test * save path * use torch * add dependency * update ignore * update training * large model * use curriculum training * add time to exp name * storage_path * restore * update training * use my key * add log message * check seed * restore callback * restore call bacl * add log message * add logging message * restore ray1.4 * length 500 * ray 100 * wandb * use tf * more levels * add callback * 10 worker * show level * no env horizon * callback result level * more call back * add diffuculty * add mroen stat * mroe stat * show levels * add callback * new * ep len 600 * fix setup * fix stepup * fix to 3.8 * update setup * parallel worker! * new exp * add callback * lateral dist * pg dataset * evaluate * modify config * align config * train single RL * update training script * 100w eval * less eval to reveal * 2000 env eval * new trianing * eval 1000 * update eval * more workers * more worker * 20 worker * dataset to database * split tool! * split dataset * try fix * train 003 * fix mapping * fix test * add waymo tqdm * utils * fix bug * fix bug * waymo * int type * 8 worker read * disable * read file * add log message * check existence * dist 0 * int * check num * suprass warning * add filter API * filter * store map false * new * ablation * filter * fix * update filyter * reanme to from * random select * add overlapping checj * fix * new training sceheme * new reward * add waymo train script * waymo different config * copy raw data * fix bug * add tqdm * update readme * waymo * pg * max lateral dist 3 * pg * crash_done instead of penalty * no crash done * gpu * update eval script * steering range penalty * evaluate * finish pg * update setup * fix bug * test * fix * add on line * train nuplan * generate sensor * udpate training * static obj * multi worker eval * filx bug * use ray for testing * eval! * filter senario * id filter * fox bug * dist = 2 * filter * eval * eval ret * ok * update training pg * test before use * store data=False * collect figures * capture pic --------- Co-authored-by: Quanyi Li <quanyi@bolei-gpu02.cs.ucla.edu>
2023-06-10 18:56:33 +01:00
parent 41c0b01f39
commit db50bca7fd
53 changed files with 2274 additions and 133 deletions
--- a/scenarionet_training/wandb_utils/our_wandb_callbacks.py
+++ b/scenarionet_training/wandb_utils/our_wandb_callbacks.py
@@ -0,0 +1,67 @@
+from ray.tune.integration.wandb import WandbLoggerCallback, _clean_log, \
+    Queue, WandbLogger
+
+
+class OurWandbLoggerCallback(WandbLoggerCallback):
+    def __init__(self, exp_name, *args, **kwargs):
+        super(OurWandbLoggerCallback, self).__init__(*args, **kwargs)
+        self.exp_name = exp_name
+
+    def log_trial_start(self, trial: "Trial"):
+        config = trial.config.copy()
+
+        config.pop("callbacks", None)  # Remove callbacks
+
+        exclude_results = self._exclude_results.copy()
+
+        # Additional excludes
+        exclude_results += self.excludes
+
+        # Log config keys on each result?
+        if not self.log_config:
+            exclude_results += ["config"]
+
+        # Fill trial ID and name
+        trial_id = trial.trial_id if trial else None
+        # trial_name = str(trial) if trial else None
+
+        # Project name for Wandb
+        wandb_project = self.project
+
+        # Grouping
+        wandb_group = self.group or trial.trainable_name if trial else None
+
+        # remove unpickleable items!
+        config = _clean_log(config)
+
+        assert trial_id is not None
+        run_name = "{}_{}".format(self.exp_name, trial_id)
+
+        wandb_init_kwargs = dict(
+            id=trial_id,
+            name=run_name,
+            resume=True,
+            reinit=True,
+            allow_val_change=True,
+            group=wandb_group,
+            project=wandb_project,
+            config=config
+        )
+        wandb_init_kwargs.update(self.kwargs)
+
+        self._trial_queues[trial] = Queue()
+        self._trial_processes[trial] = self._logger_process_cls(
+            queue=self._trial_queues[trial],
+            exclude=exclude_results,
+            to_config=self._config_results,
+            **wandb_init_kwargs
+        )
+        self._trial_processes[trial].start()
+
+    def __del__(self):
+        if self._trial_processes:
+            for v in self._trial_processes.values():
+                if hasattr(v, "close"):
+                    v.close()
+            self._trial_processes.clear()
+            self._trial_processes = {}