[RLlib][Training iteration fn] APEX conversion #22937
Conversation
… some code that needs to be deleted here, and some long running benchmarks that need to be run, but its looking good.
| @@ -19,7 +19,8 @@ | |||
| class TestAlphaStar(unittest.TestCase): | |||
| @classmethod | |||
| def setUpClass(cls): | |||
| ray.init(num_cpus=20) | |||
| # ray.init(num_cpus=20) | |||
| ray.init(local_mode=True) | |||
There was a problem hiding this comment.
whoops yeah I'll get rid of this my b.
…sult dictionaries
…led from number of replay actors:
| def setup(self, config: PartialTrainerConfigDict): | ||
| super().setup(config) | ||
| num_replay_buffer_shards = self.config["optimizer"]["num_replay_buffer_shards"] | ||
|
|
There was a problem hiding this comment.
Let's add an if-block here for _disable_execution_plan=True, such that for the execution_plan version of APEX, we don't create the learner thread and the actors twice.
rllib/agents/dqn/apex.py
Outdated
| training_intensity = int(self.config["training_intensity"] or 1) | ||
| for _ in range(training_intensity): | ||
| temp = ray.get([actor.replay.remote() for actor in self.replay_actors]) | ||
| temp = [ |
There was a problem hiding this comment.
Could this be simplified into only one for-loop?
| @@ -14,9 +15,13 @@ def asynchronous_parallel_requests( | |||
| actors: List[ActorHandle], | |||
| ray_wait_timeout_s: Optional[float] = None, | |||
| max_remote_requests_in_flight_per_actor: int = 2, | |||
| remote_fn: Optional[Callable[[Any, Optional[Any], Optional[Any]], Any]] = None, | |||
| remote_fn: Optional[ | |||
There was a problem hiding this comment.
Nice! Let's remove the Optional[...] then.
| @@ -51,6 +56,10 @@ def asynchronous_parallel_requests( | |||
| (kwargs) as **kwargs to be passed to the `remote_fn`. | |||
| E.g.: actors=[A, B], | |||
| remote_kwargs=[{...} <- **kwargs for A, {...} <- **kwargs for B]. | |||
| return_result_obj_ref_ids: If True, return the object ref IDs of the ready | |||
| results, otherwise return the actual results. | |||
| num_requests_to_launch: Number of remote requests to launch on each of the | |||
There was a problem hiding this comment.
Can we describe this more precisely? Like add "if we have not reached max_remote_requests_in_flight_per_actor yet, launch exactly num_requests_to_launch on the respective actor, regardless of the current values of remote_requests_in_flight". Something like this.
There was a problem hiding this comment.
gotcha -- yeah this is still "rate limited" by the max_remote_requests_in_flight_per_actor
|
Sorry, I know I already approved, but there are a few minor nits that we should address before merging (see my new comments). Nothing big, this is great PR! Thanks @avnishn . I'm currently running some benchmarks on Atari + APEX and will report the results here, but it looks promising. ... |
| @@ -283,6 +375,260 @@ def add_apex_metrics(result: dict) -> dict: | |||
| merged_op, workers, config, selected_workers=selected_workers | |||
| ).for_each(add_apex_metrics) | |||
|
|
|||
| def get_samples_and_store_to_replay_buffers(self): | |||
There was a problem hiding this comment.
Could we add a docstring here and in all new APEXTrainer methods?
…to apex_training_itr
…_training_itr # Conflicts: # rllib/agents/dqn/apex.py
Apart of the training iteration function overhaul.
There's probably stuff broken with the logging (some keys not getting logged right now)
I had to modify
asynchronous_parallel_requeststo add some extra functionality for convenienceAnd the bug I fixed was that if there were multiple inflight requests for a single actor that were ready in the same call to
asynchronous_parallel_requests, previously all but the last result was dropped from the return value ofasynchronous_parallel_requests. But I made the return value a dict corresponding from actor ids to lists of results, as opposed to just results, and caught this. We definitely need some standalone tests for this function.Why are these changes needed?
Related issue number
Checks
scripts/format.shto lint the changes in this PR.Tested on the tuned cartpole examples, and with minor hparam tuning (training intensity mainly) This implementation beats the execution plan implementation in runtime per number of sampled timesteps, and overall number of timesteps to optimal reward performance.