tdpeuter/2026SEL3-project-Brittle_Star_Controller_Modularity
tdpeuter/2026SEL3-project-Brittle_Star_Controller_Modularity
Archived
1
Fork 0
Code Issues Releases Activity

fix: formatted simulate script

Browse source
This commit is contained in:
Jona Reynaert 2026-04-15 18:57:14 +02:00
parent ea9ed21c9e
commit f6022cc913
1 changed files with 4 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

13
scripts/simulate.py
View file

@ -15,6 +15,7 @@ from brittle_star_project import (
) )
from brittle_star_project.environment import ArenaConfig, EnvConfig, MorphologyConfig, from_file from brittle_star_project.environment import ArenaConfig, EnvConfig, MorphologyConfig, from_file
def _flatten_obs_dict(obs_dict: dict[str, Any]) -> jnp.ndarray: def _flatten_obs_dict(obs_dict: dict[str, Any]) -> jnp.ndarray:
"""Flatten the env's observation dict into a 1D vector. """Flatten the env's observation dict into a 1D vector.
@ -35,7 +36,6 @@ def _flatten_obs_dict(obs_dict: dict[str, Any]) -> jnp.ndarray:
# A minimal policy class to load a CleanRL/Flax checkpoint and run inference. # A minimal policy class to load a CleanRL/Flax checkpoint and run inference.
class CleanRLPPOPolicy: class CleanRLPPOPolicy:
def __init__( def __init__(
self, self,
*, *,
@ -242,9 +242,7 @@ def _run_one_episode_viewer(
# Use the viewer as a context manager to avoid GLX teardown races # Use the viewer as a context manager to avoid GLX teardown races
# (e.g. GLXBadDrawable from X_GLXSwapBuffers after a window is destroyed). # (e.g. GLXBadDrawable from X_GLXSwapBuffers after a window is destroyed).
with mujoco.viewer.launch_passive(model, data) as viewer: with mujoco.viewer.launch_passive(model, data) as viewer:
step_iter = ( step_iter = range(int(max_steps)) if max_steps is not None else itertools.count()
range(int(max_steps)) if max_steps is not None else itertools.count()
)
for _step_idx in step_iter: for _step_idx in step_iter:
if not viewer.is_running(): if not viewer.is_running():
break break
@ -307,9 +305,7 @@ def parse_args() -> argparse.Namespace:
"--model", "--model",
type=str, type=str,
required=True, required=True,
help=( help=("Path to the Flax checkpoint saved by scripts/train.py (final_model.flax)."),
"Path to the Flax checkpoint saved by scripts/train.py (final_model.flax)."
),
) )
p.add_argument( p.add_argument(
"--headless", "--headless",
@ -379,8 +375,7 @@ def main() -> None:
model_path = Path(args.model) model_path = Path(args.model)
if model_path.name != "final_model.flax" or model_path.suffix != ".flax": if model_path.name != "final_model.flax" or model_path.suffix != ".flax":
raise ValueError( raise ValueError(
"Expected the training artifact 'final_model.flax', " f"Expected the training artifact 'final_model.flax', got '{model_path.name}'."
f"got '{model_path.name}'."
) )
policy = CleanRLPPOPolicy.load( policy = CleanRLPPOPolicy.load(