From 552db18da49df54a4ae55c51af8bd4d910a40a67 Mon Sep 17 00:00:00 2001 From: Theresa Eimer Date: Tue, 18 Nov 2025 14:31:20 +0100 Subject: [PATCH] fix brax context setting --- .gitignore | 4 +- carl/envs/brax/carl_brax_env.py | 3 +- changelog.md | 8 ++ examples/carl_with_sb3.py | 4 +- examples/sample_contexts_with_brax.ipynb | 4 +- pyproject.toml | 2 +- test/test_brax_env.py | 31 ++++++- test_brax.py | 102 +++++++++++++++++++++++ 8 files changed, 151 insertions(+), 7 deletions(-) create mode 100644 test_brax.py diff --git a/.gitignore b/.gitignore index 805a18b6..9acd91ce 100644 --- a/.gitignore +++ b/.gitignore @@ -33,4 +33,6 @@ generated core *.tex build -target \ No newline at end of file +target + +.DS_Store \ No newline at end of file diff --git a/carl/envs/brax/carl_brax_env.py b/carl/envs/brax/carl_brax_env.py index 6b35f84b..1f61db0c 100644 --- a/carl/envs/brax/carl_brax_env.py +++ b/carl/envs/brax/carl_brax_env.py @@ -288,8 +288,7 @@ def _update_context(self) -> None: sys = sys.replace( elasticity=sys.elasticity.at[:].set(context["elasticity"]) ) - - self.env.unwrapped.sys = sys + self.env.unwrapped._env.sys = sys def reset( self, *, seed: int | None = None, options: dict[str, Any] | None = None diff --git a/changelog.md b/changelog.md index d8e60288..d1531aa5 100644 --- a/changelog.md +++ b/changelog.md @@ -1,3 +1,11 @@ +# 1.1.2 (current) +- python version upgrade (since brax would be incompatible with 3.9) +- fix & test context propagation in brax + +# 1.1.1 +- brax version fix +- smaller bugs + # 1.1.0 - increased test coverage - smaller bug fixes diff --git a/examples/carl_with_sb3.py b/examples/carl_with_sb3.py index b606b1af..bc907213 100644 --- a/examples/carl_with_sb3.py +++ b/examples/carl_with_sb3.py @@ -8,7 +8,9 @@ from carl.context.sampler import ContextSampler # Create environment -context_distributions = [NormalFloatContextFeature("GRAVITY_X", mu=9.8, sigma=1, upper=50, lower=0)] +context_distributions = [ + NormalFloatContextFeature("GRAVITY_X", mu=9.8, sigma=1, upper=50, lower=0) +] context_sampler = ContextSampler( context_distributions=context_distributions, context_space=CARLLunarLander.get_context_space(), diff --git a/examples/sample_contexts_with_brax.ipynb b/examples/sample_contexts_with_brax.ipynb index c6473ceb..5f0bcaac 100644 --- a/examples/sample_contexts_with_brax.ipynb +++ b/examples/sample_contexts_with_brax.ipynb @@ -94,7 +94,9 @@ ], "source": [ "seed = 0\n", - "context_distributions = [NormalFloatContextFeature(\"gravity\", mu=9.8, sigma=1, upper=50, lower=0)]\n", + "context_distributions = [\n", + " NormalFloatContextFeature(\"gravity\", mu=9.8, sigma=1, upper=50, lower=0)\n", + "]\n", "context_sampler = ContextSampler(\n", " context_distributions=context_distributions,\n", " context_space=CARLBraxAnt.get_context_space(),\n", diff --git a/pyproject.toml b/pyproject.toml index 5c50361d..11de2e9c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,7 +30,7 @@ classifiers = [ "Topic :: Scientific/Engineering", "Topic :: Software Development", ] -requires-python = ">=3.9" +requires-python = ">=3.10" dependencies = ["gym", "gymnasium<1.0.0", "pygame", diff --git a/test/test_brax_env.py b/test/test_brax_env.py index 16f36de4..daca972b 100644 --- a/test/test_brax_env.py +++ b/test/test_brax_env.py @@ -1,7 +1,8 @@ import inspect import unittest -import carl.envs.gymnasium +import carl +from carl.envs.brax import CARLBraxHalfcheetah class TestBraxEnvs(unittest.TestCase): @@ -22,6 +23,34 @@ def test_envs(self): print(f"Cannot instantiate {env_name} environment.") raise e + def test_context_propagation(self): + contexts = { + 0: {"mass_torso": 20.0, "gravity": 5}, + 1: {"mass_torso": 30.0, "gravity": 15}, + } + env = CARLBraxHalfcheetah(contexts=contexts) + env.reset() + torso_idx = env.env.unwrapped._env.sys.link_names.index("torso") + + current_context = env.contexts[env.context_id] + assert env.env.unwrapped._env.sys.gravity[-1] == current_context["gravity"], ( + "Gravity not set correctly in env." + ) + assert ( + env.env.unwrapped._env.sys.link.inertia.mass[torso_idx] + == current_context["mass_torso"] + ), "Mass not set correctly in env." + + env.reset() + current_context = env.contexts[env.context_id] + assert env.env.unwrapped._env.sys.gravity[-1] == current_context["gravity"], ( + "Gravity does not change upon reset." + ) + assert ( + env.env.unwrapped._env.sys.link.inertia.mass[torso_idx] + == current_context["mass_torso"] + ), "Mass does not change upon reset." + if __name__ == "__main__": TestBraxEnvs().test_envs() diff --git a/test_brax.py b/test_brax.py new file mode 100644 index 00000000..c4b78be0 --- /dev/null +++ b/test_brax.py @@ -0,0 +1,102 @@ +from __future__ import annotations +from dataclasses import asdict +import time +import traceback + +from carl.context.selection import StaticSelector +from carl.envs import CARLBraxHalfcheetah +from gymnasium.wrappers import FlattenObservation, FilterObservation #StepAPICompatibility +import numpy as np + + +def init_carl(carl_env_fn, contexts=None, obs_context_features=None, hide_context=True, context_selector=None): + env = carl_env_fn(contexts=contexts, + obs_context_features=obs_context_features, + context_selector=context_selector) + if hide_context: + env = FlattenObservation(FilterObservation(env, filter_keys=["obs"])) + else: + env = FlattenObservation(FilterObservation(env, filter_keys=["obs", "context"])) + return env + +def main(): + '''Fixing the action sequence and changing the dynamics. Then observing the change in trjectories.''' + def print_context(eval_env): + inertia_data = asdict(eval_env.env.unwrapped._env.sys.link.inertia) + link_names = eval_env.env.unwrapped._env.sys.link_names + # print(f'link_names = {link_names}') + link_name = context_labels[0].split("_")[-1] + # print(f'link_name = {link_name}') + if link_name in link_names: + idx = link_names.index(link_name) + # inertia_data["mass"] = inertia_data["mass"].at[idx] + print(f'{link_names[idx]} = {inertia_data["mass"][idx]}') + else: + print(f'link_name={link_name} not found in link_names={link_names}.') + + carl_env_fn = CARLBraxHalfcheetah + DEFAULT_CONTEXT = carl_env_fn.get_default_context() + context_labels = ["mass_torso"] + labels = np.array(context_labels) + n_samples = 5 + rel_std = 0.25 + + context_mean = [] + for key in DEFAULT_CONTEXT.keys(): + if key in context_labels: + context_mean.append(DEFAULT_CONTEXT[key]) + context_labels = ["mass_torso"] + labels = np.array(context_labels) + + context_rel_std = rel_std + context_std = [abs(mean)*context_rel_std for mean in context_mean] + + eval_context_array = np.zeros((len(context_mean), n_samples)) + for i in range(len(context_mean)): + eval_context_array[i,:] = np.random.normal(context_mean[i], context_std[i], n_samples) + + eval_context_dict = {} + for i in range(eval_context_array.shape[1]): + eval_context_dict[i] = {0:{key:value for key,value in zip(context_labels, eval_context_array[:,i])}} + + action_seq = [] + imax = n_samples + obs_traj = [] + for i in range(imax): + contexti = eval_context_dict[i] + print(f'eval_context_dict[{i}]={eval_context_dict[i]}') + eval_env = init_carl(carl_env_fn, + contexts=contexti, + obs_context_features=context_labels, + hide_context=True, + context_selector=StaticSelector + ) + obs, _ = eval_env.reset() + print_context(eval_env) + done = False + obs_traj.append([]) + ep_rew = 0.0 + obs_traj[i].append(obs) + timestep = 0 + while not done: + if i==0: + action = eval_env.action_space.sample() + action_seq.append(action) + else: + action = action_seq[timestep] + obs, rew, term, trunc, _ = eval_env.step(action) + obs_traj[i].append(obs) + ep_rew += rew + done = term or trunc + timestep += 1 + print(f'len(obs_traj) = {len(obs_traj)}, len(obs_traj[{i}]) = {len(obs_traj[i])}, ep_rew = {ep_rew:.2f}') + + obs_traj = np.array(obs_traj) + print(f'obs_traj.shape = {obs_traj.shape}') + + for i in range(1,obs_traj.shape[0]): + obs_traj[i] = obs_traj[i] - obs_traj[0] + print(f'sum(obs_traj[{i}]) = {np.sum(np.absolute(obs_traj[i])):.4f}') + +if __name__ == "__main__": + main() \ No newline at end of file