chore: import upstream snapshot with attribution

rllib/examples/envs/classes/utils/cartpole_observations.proto

@@ -0,0 +1,8 @@
+syntax = "proto3";
+
+message CartPoleObservation {
+  double x_pos = 1;
+  double x_veloc = 2;
+  double angle_pos = 3;
+  double angle_veloc = 4;
+}

rllib/examples/envs/classes/utils/cartpole_observations_proto.py

@@ -0,0 +1,30 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: cartpole_observations.proto
+# Protobuf Python Version: 5.26.1
+"""Generated protocol buffer code."""
+from google.protobuf import (
+    descriptor as _descriptor,
+    descriptor_pool as _descriptor_pool,
+    symbol_database as _symbol_database,
+)
+from google.protobuf.internal import builder as _builder
+
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(
+    b'\n\x1b\x63\x61rtpole_observations.proto"]\n\x13\x43\x61rtPoleObservation\x12\r\n\x05x_pos\x18\x01 \x01(\x01\x12\x0f\n\x07x_veloc\x18\x02 \x01(\x01\x12\x11\n\tangle_pos\x18\x03 \x01(\x01\x12\x13\n\x0b\x61ngle_veloc\x18\x04 \x01(\x01\x62\x06proto3'  # noqa
+)
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(
+    DESCRIPTOR, "cartpole_observations_proto", _globals
+)
+if not _descriptor._USE_C_DESCRIPTORS:
+    DESCRIPTOR._loaded_options = None
+    _globals["_CARTPOLEOBSERVATION"]._serialized_start = 31
+    _globals["_CARTPOLEOBSERVATION"]._serialized_end = 124
+# @@protoc_insertion_point(module_scope)

rllib/examples/envs/classes/utils/dummy_external_client.py

@@ -0,0 +1,126 @@
+import pickle
+import socket
+import time
+
+import gymnasium as gym
+import numpy as np
+
+from ray.rllib.core import (
+    COMPONENT_RL_MODULE,
+    Columns,
+)
+from ray.rllib.env.external.rllink import (
+    RLlink,
+    get_rllink_message,
+    send_rllink_message,
+)
+from ray.rllib.env.single_agent_episode import SingleAgentEpisode
+from ray.rllib.utils.framework import try_import_torch
+from ray.rllib.utils.numpy import softmax
+
+torch, _ = try_import_torch()
+
+
+def _dummy_external_client(port: int = 5556):
+    """A dummy client that runs CartPole and acts as a testing external env."""
+
+    def _set_state(msg_body, rl_module):
+        rl_module.set_state(msg_body[COMPONENT_RL_MODULE])
+        # return msg_body[WEIGHTS_SEQ_NO]
+
+    # Connect to server.
+    while True:
+        try:
+            print(f"Trying to connect to localhost:{port} ...")
+            sock_ = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            sock_.connect(("localhost", port))
+            break
+        except ConnectionRefusedError:
+            time.sleep(5)
+
+    # Send ping-pong.
+    send_rllink_message(sock_, {"type": RLlink.PING.name})
+    msg_type, msg_body = get_rllink_message(sock_)
+    assert msg_type == RLlink.PONG
+
+    # Request config.
+    send_rllink_message(sock_, {"type": RLlink.GET_CONFIG.name})
+    msg_type, msg_body = get_rllink_message(sock_)
+    assert msg_type == RLlink.SET_CONFIG
+
+    config = pickle.loads(msg_body["config"])
+    # Create the RLModule.
+    rl_module = config.get_rl_module_spec().build()
+
+    # Request state/weights.
+    send_rllink_message(sock_, {"type": RLlink.GET_STATE.name})
+    msg_type, msg_body = get_rllink_message(sock_)
+    assert msg_type == RLlink.SET_STATE
+    _set_state(msg_body["state"], rl_module)
+
+    env_steps_per_sample = config.get_rollout_fragment_length()
+
+    # Start actual env loop.
+    env = gym.make("CartPole-v1")
+    obs, _ = env.reset()
+    episode = SingleAgentEpisode(observations=[obs])
+    episodes = [episode]
+
+    while True:
+        # Perform action inference using the RLModule.
+        logits = rl_module.forward_exploration(
+            batch={
+                Columns.OBS: torch.tensor(np.array([obs], np.float32)),
+            }
+        )[Columns.ACTION_DIST_INPUTS][
+            0
+        ].numpy()  # [0]=batch size 1
+
+        # Stochastic sample.
+        action_probs = softmax(logits)
+        action = int(np.random.choice(list(range(env.action_space.n)), p=action_probs))
+        logp = float(np.log(action_probs[action]))
+
+        # Perform the env step.
+        obs, reward, terminated, truncated, _ = env.step(action)
+
+        # Collect step data.
+        episode.add_env_step(
+            action=action,
+            reward=reward,
+            observation=obs,
+            terminated=terminated,
+            truncated=truncated,
+            extra_model_outputs={
+                Columns.ACTION_DIST_INPUTS: logits,
+                Columns.ACTION_LOGP: logp,
+            },
+        )
+
+        # We collected enough samples -> Send them to server.
+        if sum(map(len, episodes)) == env_steps_per_sample:
+            # Send the data to the server.
+            send_rllink_message(
+                sock_,
+                {
+                    "type": RLlink.EPISODES_AND_GET_STATE.name,
+                    "episodes": [e.get_state() for e in episodes],
+                    "timesteps": env_steps_per_sample,
+                },
+            )
+            # We are forced to sample on-policy. Have to wait for a response
+            # with the state (weights) in it.
+            msg_type, msg_body = get_rllink_message(sock_)
+            assert msg_type == RLlink.SET_STATE
+            _set_state(msg_body["state"], rl_module)
+
+            episodes = []
+            if not episode.is_done:
+                episode = episode.cut()
+                episodes.append(episode)
+
+        # If episode is done, reset env and create a new episode.
+        if episode.is_done:
+            obs, _ = env.reset()
+            episode = SingleAgentEpisode(observations=[obs])
+            episodes.append(episode)