From c9989fa1f3065dd2e656a8a8761d1a6277bbf844 Mon Sep 17 00:00:00 2001 From: Yarden Date: Tue, 10 Jun 2025 16:23:06 +0200 Subject: [PATCH 1/5] Fix bugs in how transitions are parserd --- ss2r/algorithms/sac/data.py | 45 ++++++++++++++++++++++++++----------- ss2r/rl/online.py | 24 +++++++++++--------- 2 files changed, 46 insertions(+), 23 deletions(-) diff --git a/ss2r/algorithms/sac/data.py b/ss2r/algorithms/sac/data.py index 690728836..7c912b42f 100644 --- a/ss2r/algorithms/sac/data.py +++ b/ss2r/algorithms/sac/data.py @@ -6,7 +6,7 @@ from brax.training import acting from brax.training.acme import running_statistics from brax.training.replay_buffers import ReplayBuffer -from brax.training.types import Params, PRNGKey +from brax.training.types import Params, PRNGKey, Transition from ss2r.algorithms.sac.go1_sac_to_onnx import convert_policy_to_onnx from ss2r.algorithms.sac.types import CollectDataFn, ReplayBufferState, float16 @@ -18,23 +18,30 @@ def get_collection_fn(cfg): if cfg.agent.data_collection.name == "step": return collect_single_step elif cfg.agent.data_collection.name == "episodic": - fn = ( - lambda env, + + def generate_episodic_unroll( + env, env_state, make_policy_fn, policy_params, key, - extra_fields: generate_unroll( + episode_length, + extra_fields, + ): + env_state, transitions = acting.generate_unroll( env, env_state, - make_policy_fn, - policy_params, + make_policy_fn(policy_params), key, - cfg.training.episode_length, + episode_length, extra_fields, ) - ) - return make_collection_fn(fn) + transitions = jax.tree.map( + lambda x: x.reshape(-1, *x.shape[2:]), transitions + ) + return env_state, transitions + + return make_collection_fn(generate_episodic_unroll) elif cfg.agent.data_collection.name == "hardware": data_collection_cfg = cfg.agent.data_collection if "Go1" in cfg.environment.task_name: @@ -45,6 +52,7 @@ def get_collection_fn(cfg): ) orchestrator = OnlineEpisodeOrchestrator( policy_translate_fn, + go1_postprocess_data, cfg.training.episode_length, data_collection_cfg.address, ) @@ -104,10 +112,6 @@ def collect_data( key, extra_fields=extra_fields, ) - if transitions.reward.ndim == 2: - transitions = jax.tree.map( - lambda x: x.reshape(-1, *x.shape[2:]), transitions - ) normalizer_params = running_statistics.update( normalizer_params, transitions.observation ) @@ -124,3 +128,18 @@ def make_go1_policy(make_policy_fn, params, cfg): del make_policy_fn proto_model = convert_policy_to_onnx(params, cfg, 12, 48) return proto_model.SerializeToString() + + +def go1_postprocess_data(raw_data, extra_fields): + observation, action, reward, next_observation, done, info = raw_data + state_extras = {x: info[x] for x in extra_fields} + policy_extras = {} + transitions = Transition( + observation=observation, + action=action, + reward=reward, + discount=1 - done, + next_observation=next_observation, + extras={"policy_extras": policy_extras, "state_extras": state_extras}, + ) + return transitions diff --git a/ss2r/rl/online.py b/ss2r/rl/online.py index fc3f4acb6..384263dd8 100644 --- a/ss2r/rl/online.py +++ b/ss2r/rl/online.py @@ -3,6 +3,7 @@ import cloudpickle as pickle import jax +import jax.numpy as jnp import zmq from brax import envs from brax.training import acting @@ -16,6 +17,7 @@ class OnlineEpisodeOrchestrator: def __init__( self, translate_policy_to_binary_fn, + data_postprocess_fn, num_steps, address="tcp://localhost:5555", ): @@ -40,6 +42,7 @@ def __init__( SSH target (e.g., 'user@host[:port]') to tunnel through. """ self._translate_policy_to_binary_fn = translate_policy_to_binary_fn + self._data_postprocess_fn = data_postprocess_fn self.num_steps = num_steps self._address = address @@ -53,26 +56,26 @@ def request_data( *, extra_fields: Sequence[str], ) -> Tuple[envs.State, Transition]: - dummy_transitions = acting.generate_unroll( + dummy_transitions = acting.actor_step( env, env_state, make_policy_fn(policy_params), key, - self.num_steps, extra_fields, )[1] - dummy_transitions = jax.tree.map(lambda x: x.squeeze(1), dummy_transitions) + dummy_transitions = jax.tree.map( + lambda x: jnp.tile(x, (self.num_steps,) + (1,) * (x.ndim - 1)), + dummy_transitions, + ) transitions = io_callback( - functools.partial(self._send_request, make_policy_fn), + functools.partial(self._send_request, make_policy_fn, extra_fields), dummy_transitions, policy_params, ordered=True, ) - state_extras = {x: transitions.extras["state_extras"][x] for x in extra_fields} - transitions.extras["state_extras"] = state_extras return env_state, transitions - def _send_request(self, make_policy_fn, policy_params): + def _send_request(self, make_policy_fn, extra_fields, policy_params): policy_bytes = self._translate_policy_to_binary_fn( make_policy_fn, policy_params ) @@ -84,6 +87,7 @@ def _send_request(self, make_policy_fn, policy_params): # Send data socket.send(pickle.dumps((policy_bytes, self.num_steps))) # Receive response - success, transitions = pickle.loads(socket.recv()) - if success: - return transitions + raw_data = pickle.loads(socket.recv()) + transitions = self._data_postprocess_fn(raw_data, extra_fields) + print(f"Received {len(transitions.reward)} transitions...") + return transitions From 971d3be5a6fcc5734fe96ed7222750bac48caee8 Mon Sep 17 00:00:00 2001 From: Yarden Date: Tue, 10 Jun 2025 17:24:54 +0200 Subject: [PATCH 2/5] Implement lazy pirate pattern --- ss2r/rl/online.py | 30 ++++++++++++++++++++++++------ 1 file changed, 24 insertions(+), 6 deletions(-) diff --git a/ss2r/rl/online.py b/ss2r/rl/online.py index 384263dd8..5a9b55fb6 100644 --- a/ss2r/rl/online.py +++ b/ss2r/rl/online.py @@ -12,6 +12,9 @@ from ss2r.rl.types import MakePolicyFn +_REQUEST_TIMEOUT = 120000 +_REQUEST_RETRIES = 120 + class OnlineEpisodeOrchestrator: def __init__( @@ -76,18 +79,33 @@ def request_data( return env_state, transitions def _send_request(self, make_policy_fn, extra_fields, policy_params): + """Implements a lazy pirate client reliability pattern""" policy_bytes = self._translate_policy_to_binary_fn( make_policy_fn, policy_params ) with zmq.Context() as ctx: - with ctx.socket(zmq.REQ) as socket: - socket.connect(self._address) - while True: - print("Requesting data...") - # Send data - socket.send(pickle.dumps((policy_bytes, self.num_steps))) + socket = ctx.socket(zmq.REQ) + socket.connect(self._address) + retries_left = _REQUEST_RETRIES + print("Requesting data...") + # Send data + socket.send(pickle.dumps((policy_bytes, self.num_steps))) + while True: + if (socket.poll(_REQUEST_TIMEOUT) & zmq.POLLIN) != 0: # Receive response raw_data = pickle.loads(socket.recv()) transitions = self._data_postprocess_fn(raw_data, extra_fields) print(f"Received {len(transitions.reward)} transitions...") return transitions + else: + retries_left -= 1 + print(f"Request timed out, {retries_left} retries left...") + socket.setsockopt(zmq.LINGER, 0) + socket.close() + if retries_left == 0: + raise RuntimeError("Request timed out.") + print("Retrying...") + socket = ctx.socket(zmq.REQ) + socket.connect(self._address) + print("Requesting data...") + socket.send(pickle.dumps((policy_bytes, self.num_steps))) From c198cdc1b36ac10d569d90806550f9790b5f763c Mon Sep 17 00:00:00 2001 From: Yarden Date: Tue, 10 Jun 2025 17:56:55 +0200 Subject: [PATCH 3/5] Default --- ss2r/algorithms/sac/data.py | 2 +- ss2r/rl/online.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/ss2r/algorithms/sac/data.py b/ss2r/algorithms/sac/data.py index 7c912b42f..48f7ab1d6 100644 --- a/ss2r/algorithms/sac/data.py +++ b/ss2r/algorithms/sac/data.py @@ -52,8 +52,8 @@ def generate_episodic_unroll( ) orchestrator = OnlineEpisodeOrchestrator( policy_translate_fn, - go1_postprocess_data, cfg.training.episode_length, + go1_postprocess_data, data_collection_cfg.address, ) return make_collection_fn(orchestrator.request_data) diff --git a/ss2r/rl/online.py b/ss2r/rl/online.py index 5a9b55fb6..4b3669663 100644 --- a/ss2r/rl/online.py +++ b/ss2r/rl/online.py @@ -20,8 +20,8 @@ class OnlineEpisodeOrchestrator: def __init__( self, translate_policy_to_binary_fn, - data_postprocess_fn, num_steps, + data_postprocess_fn=lambda x, y: x, address="tcp://localhost:5555", ): """Orchestrator for requesting episodes over ZMQ, with optional SSH reverse tunnel. From 140a72c6bbbb69cc8e490d9a89264d840cd2e65f Mon Sep 17 00:00:00 2001 From: Yarden Date: Tue, 10 Jun 2025 18:03:19 +0200 Subject: [PATCH 4/5] fix bugs --- tests/test_online_orchestrator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_online_orchestrator.py b/tests/test_online_orchestrator.py index 3a04b4bd0..619d611c3 100644 --- a/tests/test_online_orchestrator.py +++ b/tests/test_online_orchestrator.py @@ -78,6 +78,6 @@ def test_send_request(): address="tcp://localhost:5555", ) policy = DummyPolicy(param=123) - transitions = orchestrator._send_request(lambda x: x, policy) + transitions = orchestrator._send_request(lambda x: x, policy, ()) assert transitions.observation["privileged_state"].shape == (100, 123) assert transitions.observation["state"].shape == (100, 48) From d025e84a7444b9c9fe3f5e316437f00d0d4af141 Mon Sep 17 00:00:00 2001 From: Yarden Date: Tue, 10 Jun 2025 18:19:24 +0200 Subject: [PATCH 5/5] remove test --- tests/test_online_orchestrator.py | 83 ------------------------------- 1 file changed, 83 deletions(-) delete mode 100644 tests/test_online_orchestrator.py diff --git a/tests/test_online_orchestrator.py b/tests/test_online_orchestrator.py deleted file mode 100644 index 619d611c3..000000000 --- a/tests/test_online_orchestrator.py +++ /dev/null @@ -1,83 +0,0 @@ -import threading -import time - -import cloudpickle as pickle -import jax -import numpy as np -import pytest -import zmq -from brax.training.types import Transition - -from ss2r.rl.online import OnlineEpisodeOrchestrator - - -# Dummy ZMQ server -def dummy_server(address="tcp://*:5555", delay=0.1): - context = zmq.Context() - socket = context.socket(zmq.REP) - socket.bind(address) - try: - while True: - message = socket.recv() - _, num_steps = pickle.loads(message) - transitions = Transition( - observation={ - "privileged_state": np.random.rand(num_steps, 123), - "state": np.random.rand(num_steps, 48), - }, - action=np.random.rand(num_steps, 12), - reward=np.random.rand(num_steps), - discount=np.random.rand(num_steps), - next_observation={ - "privileged_state": np.random.rand(num_steps, 123), - "state": np.random.rand(num_steps, 48), - }, - extras={ - "policy_extras": {}, - "state_extras": { - "cost": np.random.rand(num_steps), - "truncation": np.random.rand(num_steps), - }, - }, - ) - transitions = jax.tree.map(lambda x: x.astype(np.float32), transitions) - response = (True, transitions) - time.sleep(delay) - socket.send(pickle.dumps(response)) - except zmq.ZMQError: - pass - finally: - socket.close() - context.term() - - -# Dummy policy and translation function -class DummyPolicy: - def __init__(self, param): - self.param = param - - -def dummy_translate_policy_to_binary_fn(policy_fn, policy): - return pickle.dumps(policy) - - -# Fixture for server setup -@pytest.fixture(scope="module", autouse=True) -def start_dummy_server(): - thread = threading.Thread(target=dummy_server, daemon=True) - thread.start() - time.sleep(0.5) # Ensure server is ready - yield - # Server thread exits automatically due to daemon=True - - -def test_send_request(): - orchestrator = OnlineEpisodeOrchestrator( - translate_policy_to_binary_fn=dummy_translate_policy_to_binary_fn, - num_steps=100, - address="tcp://localhost:5555", - ) - policy = DummyPolicy(param=123) - transitions = orchestrator._send_request(lambda x: x, policy, ()) - assert transitions.observation["privileged_state"].shape == (100, 123) - assert transitions.observation["state"].shape == (100, 48)