[RLlib] External Env support for new API stack: New example script and custom tcp-capable EnvRunner. #49033

sven1977 · 2024-12-03T10:37:18Z

External Env support for new API stack:

Custom tcp-capable EnvRunner.
Example script with dummy CartPole TCP client acting as the external env.

Why are these changes needed?

Related issue number

Checks

I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
I've run scripts/format.sh to lint the changes in this PR.
I've included any doc changes needed for https://docs.ray.io/en/master/.
- I've added any new APIs to the API Reference. For example, if I added a
  method in Tune, I've added it in doc/source/tune/api/ under the
  corresponding .rst file.
I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
Testing Strategy
- Unit tests
- Release tests
- This PR is not tested :(

…rnal_env_support_new_api_stack

simonsays1980

LGTM. A couple of questions in regard to the basic design. Great PR!!

simonsays1980 · 2024-12-03T16:17:15Z

doc/source/rllib/rllib-examples.rst

 - `Env Rendering and Recording <https://github.com/ray-project/ray/blob/master/rllib/examples/envs/env_rendering_and_recording.py>`__:
   Illustrates environment rendering and recording setups within RLlib, capturing visual outputs for later review (ex. on WandB), which is essential
   for tracking agent behavior in training.

- `Env with Protobuf Observations <https://github.com/ray-project/ray/blob/master/rllib/examples/envs/env_with_protobuf_observations.py>`__:
+- `Env with Protobuf Observations <https://github.com/ray-project/ray/blob/master/rllib/examples/envs/env_w_protobuf_observations.py>`__:


Awesome that we have now this example also in new API stack! Great work!

simonsays1980 · 2024-12-03T16:19:13Z

rllib/algorithms/ppo/torch/ppo_torch_learner.py

+        # TODO (sven): We should ideally do this in the LearnerConnector (separation of
+        #  concerns: Only do things on the EnvRunners that are required for computing
+        #  actions, do NOT do anything on the EnvRunners that's only required for a
+        #   training update).


I thought we already moved that from EnvRunners into a connector?! Yes, at best it would resit in there.

It is in a connector, but in the env-to-module connector :D on the EnvRunner side, where it should ideally be run on the learner side.

simonsays1980 · 2024-12-03T16:20:48Z

rllib/core/rl_module/torch/torch_rl_module.py

@@ -177,7 +177,11 @@ def forward(self, batch: Dict[str, Any], **kwargs) -> Dict[str, Any]:
        training sample collection (w/ exploration behavior).
        `_forward_train()` to define the forward pass prior to loss computation.
        """
-        return self.forward_train(batch, **kwargs)
+        # TODO (sven): Experimental to make ONNX exported models work.
+        if self.config.inference_only:


We should think about how inference-only should be integrated such that ONNX works well. Maybe it needs a different integration in the modules.

simonsays1980 · 2024-12-03T16:23:04Z

rllib/env/multi_agent_env_runner.py

-        # will not have an RLModule, but might still be usable with random actions.
-        except NotImplementedError:
-            self.module = None
+        self.make_module()


Nice, now env runners work it similarly.

simonsays1980 · 2024-12-03T16:32:41Z

rllib/env/tcp_client_inference_env_runner.py

+
+    This implementation assumes:
+    - Only one external client ever connects to this env runner.
+    - The external client performs inference locally through an ONNX model. Thus,


An alternative would be to define an EnvRunner that can execute an external environment inside a thread - so users just define the EnvExecutor that wraps it and defines protocols to start the environment and communicate with it (sending in actions and receiving observations and rewards) and the EnvRunner can then execute. It would keep the RLModule like the other EnvRunner classes we have. Do we have any cases where this would not work? In my C++ external environment this does work quite well. What would be more pro and contras?

simonsays1980 · 2024-12-03T16:58:03Z

rllib/env/tcp_client_inference_env_runner.py

+            self.client_socket,
+            {
+                "type": MessageTypes.SET_CONFIG.name,
+                "env_steps_per_sample": self.config.get_rollout_fragment_length(


Ah okay, so per sample each external environment does only sample enough for the next batch.

Correct. This is purely for PPO for now. For IMPALA/APPO or offline algos, this doesn't really matter. The most difficult thing is making this work for PPO, b/c in this purely on-policy case, client and server have to meet exactly in the middle, at the handover of the samples and the receiving of the new weights.

simonsays1980 · 2024-12-03T16:58:32Z

rllib/env/tcp_client_inference_env_runner.py

+                "env_steps_per_sample": self.config.get_rollout_fragment_length(
+                    worker_index=self.worker_index
+                ),
+                "force_on_policy": True,


Alright, it could be off-policy if a user wants it. Cool!

Correct, then the protocol will change slightly, but not much. Clients can send data any time they want and - from time to time - get state/weights back on some of these requests.

simonsays1980 · 2024-12-03T17:04:19Z

rllib/env/tcp_client_inference_env_runner.py

+    client_thread.start()
+
+    # Query the EnvRunner.
+    episodes = env_runner.sample()


Nice! But question: In case I want to run with tune I could simply pass in: TcpClientInferenceEnvRunner into the env_runner_class attribute?

Removed this code. It was an early testing leftover.

simonsays1980 · 2024-12-03T17:06:34Z

rllib/examples/envs/env_connecting_to_rllib_w_tcp_client.py

+collects episodes of experiences, and sends these (in bulk) back to RLlib for training.
+Also, from time to time, the updated model weights have to be sent from RLlib (server)
+back to the connected clients.
+Note that RLlib's new API stack does not yet support individual action requests, where


Why is this? Only because it is built in the TcpEnvRunner correct?

Correct, this protocol (remote action inference on the RLlib server) is not supported yet. We need just another EnvRunner for that or build this logic into the existing one.

simonsays1980 · 2024-12-03T17:08:11Z

rllib/examples/envs/env_connecting_to_rllib_w_tcp_client.py

+        )
+        .env_runners(
+            # Point RLlib to the custom EnvRunner to be used here.
+            env_runner_cls=TcpClientInferenceEnvRunner,


Yup, question answered :)

doc/source/rllib/package_ref/env/utils.rst

Signed-off-by: Sven Mika <sven@anyscale.io>

…d custom tcp-capable EnvRunner. (ray-project#49033)

…d custom tcp-capable EnvRunner. (ray-project#49033) Signed-off-by: Connor Sanders <connor@elastiflow.com>

…d custom tcp-capable EnvRunner. (ray-project#49033) Signed-off-by: hjiang <dentinyhao@gmail.com>

…d custom tcp-capable EnvRunner. (ray-project#49033) Signed-off-by: ujjawal-khare <ujjawal.khare@dream11.com>

sven1977 added 5 commits November 27, 2024 16:58

wip

941743f

wip

8fe0300

wip

2daac5b

Merge branch 'master' of https://github.com/ray-project/ray into exte…

107ade2

…rnal_env_support_new_api_stack

wip

aa4ca26

sven1977 requested review from maxpumperla, simonsays1980 and a team as code owners December 3, 2024 10:37

sven1977 added 2 commits December 3, 2024 12:39

wip

80f7e3a

LINT

2de43b9

sven1977 enabled auto-merge (squash) December 3, 2024 13:18

github-actions bot added the go add ONLY when ready to merge, run all tests label Dec 3, 2024

wip

8a5af9c

github-actions bot disabled auto-merge December 3, 2024 14:26

sven1977 added 2 commits December 3, 2024 15:27

wip

edbb110

wip

5986a29

simonsays1980 approved these changes Dec 3, 2024

View reviewed changes

sven1977 commented Dec 3, 2024

View reviewed changes

doc/source/rllib/package_ref/env/utils.rst Outdated Show resolved Hide resolved

Apply suggestions from code review

c618584

Signed-off-by: Sven Mika <sven@anyscale.io>

sven1977 enabled auto-merge (squash) December 3, 2024 18:01

sven1977 merged commit 7c51a74 into ray-project:master Dec 3, 2024
5 of 6 checks passed

simonsays1980 pushed a commit to simonsays1980/ray that referenced this pull request Dec 3, 2024

[RLlib] External Env support for new API stack: New example script an…

8d19e4d

…d custom tcp-capable EnvRunner. (ray-project#49033)

jecsand838 pushed a commit to jecsand838/ray that referenced this pull request Dec 4, 2024

[RLlib] External Env support for new API stack: New example script an…

53bce12

…d custom tcp-capable EnvRunner. (ray-project#49033) Signed-off-by: Connor Sanders <connor@elastiflow.com>

dentiny pushed a commit to dentiny/ray that referenced this pull request Dec 7, 2024

[RLlib] External Env support for new API stack: New example script an…

8886121

…d custom tcp-capable EnvRunner. (ray-project#49033) Signed-off-by: hjiang <dentinyhao@gmail.com>

ujjawal-khare pushed a commit to ujjawal-khare-27/ray that referenced this pull request Dec 17, 2024

[RLlib] External Env support for new API stack: New example script an…

e564331

…d custom tcp-capable EnvRunner. (ray-project#49033) Signed-off-by: ujjawal-khare <ujjawal.khare@dream11.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[RLlib] External Env support for new API stack: New example script and custom tcp-capable EnvRunner. #49033

[RLlib] External Env support for new API stack: New example script and custom tcp-capable EnvRunner. #49033

sven1977 commented Dec 3, 2024 •

edited

Loading

simonsays1980 left a comment

simonsays1980 Dec 3, 2024

simonsays1980 Dec 3, 2024

sven1977 Dec 3, 2024

simonsays1980 Dec 3, 2024

simonsays1980 Dec 3, 2024

simonsays1980 Dec 3, 2024

simonsays1980 Dec 3, 2024

sven1977 Dec 3, 2024

simonsays1980 Dec 3, 2024

sven1977 Dec 3, 2024

simonsays1980 Dec 3, 2024

sven1977 Dec 3, 2024

simonsays1980 Dec 3, 2024

sven1977 Dec 3, 2024

simonsays1980 Dec 3, 2024

[RLlib] External Env support for new API stack: New example script and custom tcp-capable EnvRunner. #49033

[RLlib] External Env support for new API stack: New example script and custom tcp-capable EnvRunner. #49033

Conversation

sven1977 commented Dec 3, 2024 • edited Loading

Why are these changes needed?

Related issue number

Checks

simonsays1980 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sven1977 commented Dec 3, 2024 •

edited

Loading