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>

scenarionet_training/wandb_utils/__init__.py

@@ -0,0 +1,3 @@
+import os
+
+WANDB_KEY_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "wandb_api_key_file.txt")

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 = {}

scenarionet_training/wandb_utils/our_wandb_callbacks_ray100.py

@@ -0,0 +1,80 @@
+from multiprocessing import Queue
+
+from ray.tune.integration.wandb import WandbLogger, _clean_log, _set_api_key
+
+
+class OurWandbLogger(WandbLogger):
+    def __init__(self, config, logdir, trial):
+        self.exp_name = config["logger_config"]["wandb"].pop("exp_name")
+        super(OurWandbLogger, self).__init__(config, logdir, trial)
+
+    def _init(self):
+
+        config = self.config.copy()
+
+        config.pop("callbacks", None)  # Remove callbacks
+
+        try:
+            if config.get("logger_config", {}).get("wandb"):
+                logger_config = config.pop("logger_config")
+                wandb_config = logger_config.get("wandb").copy()
+            else:
+                wandb_config = config.pop("wandb").copy()
+        except KeyError:
+            raise ValueError(
+                "Wandb logger specified but no configuration has been passed. "
+                "Make sure to include a `wandb` key in your `config` dict "
+                "containing at least a `project` specification.")
+
+        _set_api_key(wandb_config)
+
+        exclude_results = self._exclude_results.copy()
+
+        # Additional excludes
+        additional_excludes = wandb_config.pop("excludes", [])
+        exclude_results += additional_excludes
+
+        # Log config keys on each result?
+        log_config = wandb_config.pop("log_config", False)
+        if not log_config:
+            exclude_results += ["config"]
+
+        # Fill trial ID and name
+        trial_id = self.trial.trial_id if self.trial else None
+        trial_name = str(self.trial) if self.trial else None
+
+        # Project name for Wandb
+        try:
+            wandb_project = wandb_config.pop("project")
+        except KeyError:
+            raise ValueError(
+                "You need to specify a `project` in your wandb `config` dict.")
+
+        # Grouping
+        wandb_group = wandb_config.pop(
+            "group", self.trial.trainable_name if self.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(wandb_config)
+
+        self._queue = Queue()
+        self._wandb = self._logger_process_cls(
+            queue=self._queue,
+            exclude=exclude_results,
+            to_config=self._config_results,
+            **wandb_init_kwargs)
+        self._wandb.start()

scenarionet_training/wandb_utils/test_wandb.py

@@ -0,0 +1,38 @@
+"""
+Procedure to use wandb:
+
+1. Logup in wandb: https://wandb.ai/
+2. Get the API key in personal setting
+3. Store API key (a string)to some file as: ~/wandb_api_key_file.txt
+4. Install wandb: pip install wandb
+5. Fill the "wandb_key_file", "wandb_project" keys in our train function.
+
+Note1: You don't need to specify who own "wandb_project", for example, in team "drivingforce"'s project
+"representation", you only need to fill wandb_project="representation"
+
+Note2: In wanbd, there are "team name", "project name", "group name" and "trial_name". We only need to care
+"team name" and "project name". The "team name" is set to "drivingforce" by default. You can also use None to
+log result to your personal domain. The "group name" of the experiment is exactly the "exp_name" in our context, like
+"0304_train_ppo" or so.
+
+Note3: It would be great to change the x-axis in wandb website to "timesteps_total".
+
+Peng Zhenghao, 20210402
+"""
+from ray import tune
+
+from scenarionet_training.train_utils.utils import train
+
+if __name__ == "__main__":
+    config = dict(env="CartPole-v0", num_workers=0, lr=tune.grid_search([1e-2, 1e-4]))
+    train(
+        "PPO",
+        exp_name="test_wandb",
+        stop=10000,
+        config=config,
+        custom_callback=False,
+        test_mode=False,
+        local_mode=False,
+        wandb_project="TEST",
+        wandb_team="drivingforce"  # drivingforce is set to default. Use None to log to your personal domain!
+    )

scenarionet_training/wandb_utils/wandb_api_key_file.txt

@@ -0,0 +1 @@
+132a8add578bdaeea5ab7a4942f35f2a17742df2