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

Merge pull request #57 from SELab-3-2026/feat/video-simulation

Browse source 
Simulation visualisations for poster
This commit is contained in:
Tibo De Peuter 2026-05-20 15:02:18 +02:00 • committed by GitHub
commit 54726e4e08
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 639 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

1
.gitignore vendored
View file

@ -6,6 +6,7 @@ outputs/
multirun/
metrics/
adjacency_debug.txt
vids/
# Python-generated files
__pycache__/

4
configs/simulation/default.yaml
View file

@ -21,6 +21,10 @@ video_output_path: null
# Camera ID to use for video recording (1 is usually the close-up camera)
camera_id: 1
video_width: 640
video_height: 80
video_fps: 60
# Optional override for the metadata YAML file path.
# If null, the script looks for `<model_name>_metadata.yaml` alongside the model_path.
metadata_path: null

13
docs/api/simulation.md
View file

@ -38,4 +38,17 @@ uv run scripts/simulate.py \
Videos and evaluation metadata are stored in timestamped folders alongside the model:
`runs/your_run/final_model_evaluations/eval_<timestamp>/simulation.mp4`
### Top-Down and Follow Cameras
Using the following script, you can render a top-down and follow camera view for multiple models at once:
```bash
uv run scripts/poster_visualisations/render_poster_videos.py \
runs/final-models/centralized/.../final_model.flax \
runs/final-models/fully-connected/.../final_model.flax \
runs/final-models/ring/.../final_model.flax \
--max-steps 10000 --width 640 --height 480 --fps 60 \
--output-root vids/poster/
```
For batch evaluation and cross-model comparison, see the **[Evaluation Guide](./evaluation.md)**.

150
scripts/poster_visualisations/render_poster_videos.py Normal file
View file

@ -0,0 +1,150 @@
from __future__ import annotations
import argparse
from pathlib import Path
from brittle_star_project.evaluation.checkpoint import load_metadata, metadata_to_configs
from brittle_star_project.evaluation.eval_env_builder import build_eval_env
from brittle_star_project.evaluation.video import record_episode_multi_camera
_ARCH_DIR_MAP = {
"CENTRALIZED": "centralized",
"FULLY_CONNECTED": "fully-connected",
"RING": "ring",
"SEGMENT": "segment",
}
ROBOT_COLOR_MAP = {
"CENTRALIZED": "#0D567C", # Blue
"FULLY_CONNECTED": "#8C0E0F", # Reddish
"RING": "#FCB304", # Pale Yellow
}
def _arch_dir(name: str) -> str:
return _ARCH_DIR_MAP.get(name, name.lower())
def _resolve_overrides(overrides: list[str], count: int) -> list[str | None]:
if not overrides:
return [None] * count
if len(overrides) == 1 and count > 1:
return overrides * count
if len(overrides) != count:
raise ValueError("morphology overrides must match the number of models")
return overrides
def main() -> None:
parser = argparse.ArgumentParser(description="Render top-down and follow videos for poster.")
parser.add_argument("models", nargs="+", help="Paths to .flax checkpoints")
parser.add_argument(
"--morphology-override",
action="append",
default=[],
help="Override morphology YAML path (repeat to match models)",
)
parser.add_argument("--output-root", default="vids/poster")
parser.add_argument("--seed", type=int, default=0)
parser.add_argument("--max-steps", type=int, default=5000)
parser.add_argument("--topdown-camera", type=int, default=0)
parser.add_argument("--follow-camera", type=int, default=1)
parser.add_argument("--topdown-camera-x", type=float, default=-3.0)
parser.add_argument("--topdown-camera-y", type=float, default=0.0)
parser.add_argument("--topdown-camera-z", type=float, default=4.5)
parser.add_argument("--topdown-camera-fovy", type=float, default=None)
parser.add_argument("--target-x", type=float, default=-6.0)
parser.add_argument("--target-y", type=float, default=0.0)
parser.add_argument("--width", type=int, default=2160)
parser.add_argument("--height", type=int, default=960)
parser.add_argument("--fps", type=int, default=60)
parser.add_argument(
"--robot-color",
default="#2B4162",
help="Hex color for the brittle star robot",
)
args = parser.parse_args()
if (args.target_x is None) != (args.target_y is None):
raise ValueError("target-x and target-y must be provided together")
target_xy = None
if args.target_x is not None:
target_xy = (float(args.target_x), float(args.target_y))
camera_fovy = None
if args.topdown_camera_fovy is not None:
camera_fovy = {args.topdown_camera: float(args.topdown_camera_fovy)}
camera_x = None
if args.topdown_camera_x is not None:
camera_x = {args.topdown_camera: float(args.topdown_camera_x)}
camera_y = None
if args.topdown_camera_y is not None:
camera_y = {args.topdown_camera: float(args.topdown_camera_y)}
camera_z = None
if args.topdown_camera_z is not None:
camera_z = {args.topdown_camera: float(args.topdown_camera_z)}
camera_xyz = (camera_x, camera_y, camera_z)
overrides = _resolve_overrides(args.morphology_override, len(args.models))
output_root = Path(args.output_root)
for model_path_str, override in zip(args.models, overrides):
model_path = Path(model_path_str)
metadata = load_metadata(model_path, None)
training = metadata_to_configs(metadata)
bundle = build_eval_env(
model_path=model_path,
training=training,
metadata=metadata,
morphology_override_path=override,
)
arch_dir = _arch_dir(bundle.architecture)
arms_dir = f"{bundle.num_active_arms}arms"
out_dir = output_root / arms_dir / arch_dir
out_dir.mkdir(parents=True, exist_ok=True)
output_paths = {
args.topdown_camera: out_dir / "topdown.mp4",
args.follow_camera: out_dir / "follow.mp4",
}
print(bundle.architecture)
color = ROBOT_COLOR_MAP.get(bundle.architecture, args.robot_color)
result = record_episode_multi_camera(
env=bundle.env,
policy=bundle.policy,
seed=args.seed,
max_steps=args.max_steps,
action_low=bundle.action_low,
action_high=bundle.action_high,
action_mask=bundle.action_mask,
output_paths=output_paths,
camera_ids=[args.topdown_camera, args.follow_camera],
camera_fovy=camera_fovy,
camera_xyz=camera_xyz,
target_xy=target_xy,
robot_color=color,
width=args.width,
height=args.height,
fps=args.fps,
)
final_dist = "n/a" if result.final_xy_dist is None else f"{result.final_xy_dist:.3f}"
print(
f"{arms_dir}/{arch_dir}: return={result.return_:.6f}, len={result.length}, "
f"target_reached={result.reached_target}, final_xy_dist={final_dist}"
)
bundle.env.close()
if __name__ == "__main__":
main()

10
scripts/poster_visualisations/render_poster_videos.sh Executable file
View file

@ -0,0 +1,10 @@
#!/usr/bin/env bash
# Multi-camera renders per model -> vids/poster/{arms}arms/{arch}/topdown.mp4 + follow.mp4
path=$1
uv run scripts/poster_visualisations/render_poster_videos.py \
"$path"/centralized.flax \
"$path"/fully-connected.flax \
"$path"/ring.flax \

241
scripts/poster_visualisations/render_static_path_image.py Normal file
View file

@ -0,0 +1,241 @@
from __future__ import annotations
import argparse
from pathlib import Path
import numpy as np
from brittle_star_project.evaluation.checkpoint import load_metadata, metadata_to_configs
from brittle_star_project.evaluation.eval_env_builder import build_eval_env
from brittle_star_project.evaluation.rollout import (
_get_observations,
_maybe_clip_action,
_target_reached,
)
from brittle_star_project.evaluation.video import (
_apply_camera_overrides,
_ensure_offscreen_size,
hex_to_rgba,
)
ROBOT_COLOR_MAP = {
"CENTRALIZED": "#0D567C", # Blue
"FULLY_CONNECTED": "#8C0E0F", # Reddish
"RING": "#FCB304", # Pale Yellow
}
def _enum_value(enum_obj, *names: str) -> int:
for name in names:
if hasattr(enum_obj, name):
return int(getattr(enum_obj, name))
raise AttributeError(f"Could not find any of {names!r} on {enum_obj!r}")
def _append_sphere(scene, mujoco, center: np.ndarray, radius: float, rgba: np.ndarray) -> None:
geom = scene.geoms[scene.ngeom]
mujoco.mjv_initGeom(
geom,
mujoco.mjtGeom.mjGEOM_SPHERE,
np.asarray([radius, 0.0, 0.0], dtype=np.float32),
center,
np.eye(3, dtype=np.float32).reshape(-1),
rgba,
)
scene.ngeom += 1
def main() -> None:
parser = argparse.ArgumentParser(description="Render a static path image from a rollout.")
parser.add_argument("model", help="Path to .flax checkpoint")
parser.add_argument("--morphology-override", default=None)
parser.add_argument("--output-path", required=True)
parser.add_argument("--seed", type=int, default=0)
parser.add_argument("--max-steps", type=int, default=5000)
parser.add_argument("--body-name", default="BrittleStarMorphology/central_disk")
parser.add_argument("--camera-id", type=int, default=0)
parser.add_argument("--camera-x", type=float, default=-3.0)
parser.add_argument("--camera-y", type=float, default=0.0)
parser.add_argument("--camera-z", type=float, default=4.5)
parser.add_argument("--camera-fovy", type=float, default=None)
parser.add_argument("--target-x", type=float, default=-6.0)
parser.add_argument("--target-y", type=float, default=0.0)
parser.add_argument("--width", type=int, default=2160)
parser.add_argument("--height", type=int, default=960)
parser.add_argument("--frame-stride", type=int, default=15)
parser.add_argument(
"--path-color", default="#FA9F42"
) # ring = #888888, centralized = #2B4162, fully connected = FA9F42
parser.add_argument(
"--robot-color",
default="#FA9F42",
help="Hex color for brittle star robot (e.g. #ff0000)",
)
args = parser.parse_args()
model_path = Path(args.model)
metadata = load_metadata(model_path, None)
training = metadata_to_configs(metadata)
bundle = build_eval_env(
model_path=model_path,
training=training,
metadata=metadata,
morphology_override_path=args.morphology_override,
)
if (args.target_x is None) != (args.target_y is None):
raise ValueError("target-x and target-y must be provided together")
target_xy = None
if args.target_x is not None:
target_xy = (float(args.target_x), float(args.target_y))
try:
import imageio
import mujoco
except ImportError as e:
raise ImportError(
"Static image rendering requires 'mujoco' and 'imageio'. "
"Please install the evaluation dependencies: `uv pip install .[evaluation]`"
) from e
reset_kwargs = {}
if target_xy is not None:
reset_kwargs["target_position"] = (target_xy[0], target_xy[1], 0.0)
state = bundle.env.reset(seed=args.seed, **reset_kwargs)
model = state.mj_model
data = state.mj_data
_apply_camera_overrides(
model,
camera_fovy={args.camera_id: float(args.camera_fovy)}
if args.camera_fovy is not None
else None,
camera_xyz=(
{args.camera_id: float(args.camera_x)} if args.camera_x is not None else None,
{args.camera_id: float(args.camera_y)} if args.camera_y is not None else None,
{args.camera_id: float(args.camera_z)} if args.camera_z is not None else None,
),
)
_ensure_offscreen_size(model, args.width, args.height)
body_id = mujoco.mj_name2id(
model, mujoco.mjtObj.mjOBJ_BODY, "BrittleStarMorphology/central_disk"
)
robot_rgba = hex_to_rgba(ROBOT_COLOR_MAP.get(bundle.architecture, args.robot_color), 1.0)
# Optionally override robot color by recoloring geoms belonging to the robot's body subtree.
if args.robot_color is not None:
# Collect body IDs in the subtree rooted at `body_id` by walking parent links.
nbody = int(model.nbody)
body_parent = model.body_parentid
robot_body_ids = set([int(body_id)])
for i in range(1, nbody):
cur = int(i)
# walk up until root (0) or until we hit the robot root
while cur not in (-1, 0, int(body_id)):
cur = int(body_parent[cur])
if cur == int(body_id):
robot_body_ids.add(i)
# Recolor geoms whose body id is in the robot subtree
for g in range(int(model.ngeom)):
if int(model.geom_bodyid[g]) in robot_body_ids:
model.geom_rgba[g][:] = robot_rgba
# Optionally override robot color by recoloring geoms belonging to the robot's body subtree.
if args.robot_color is not None:
# Collect body IDs in the subtree rooted at `body_id` by walking parent links.
nbody = int(model.nbody)
body_parent = model.body_parentid
robot_body_ids = set([int(body_id)])
for i in range(1, nbody):
cur = int(i)
# walk up until root (0) or until we hit the robot root
while cur not in (-1, 0, int(body_id)):
cur = int(body_parent[cur])
if cur == int(body_id):
robot_body_ids.add(i)
# Recolor geoms whose body id is in the robot subtree
for g in range(int(model.ngeom)):
if int(model.geom_bodyid[g]) in robot_body_ids:
model.geom_rgba[g][:] = robot_rgba
positions = []
observations = _get_observations(state)
for _ in range(int(args.max_steps)):
positions.append(np.asarray(data.xpos[body_id], dtype=np.float32))
obs_dict = observations or {}
action = bundle.policy.act(observations=obs_dict)
if bundle.action_mask is not None:
action = action[bundle.action_mask]
action = _maybe_clip_action(action, bundle.action_low, bundle.action_high)
state = bundle.env.step(state=state, action=action)
data = state.mj_data
observations = _get_observations(state)
if _target_reached(state=state):
break
positions_arr = np.vstack(positions)
if len(positions_arr) < 2:
raise ValueError("Need at least two rollout positions to render a path")
path_points = positions_arr.copy()
path_points[:, 2] -= 0.02
path_step = max(1, int(args.frame_stride))
path_points_visible = path_points[::path_step]
path_rgba = hex_to_rgba(ROBOT_COLOR_MAP.get(bundle.architecture, args.path_color), 0.92)
ctx = mujoco.GLContext(args.width, args.height)
ctx.make_current()
try:
catmask = _enum_value(mujoco.mjtCatBit, "mjCAT_ALL")
camera_type = _enum_value(mujoco.mjtCamera, "mjCAMERA_FIXED")
font_scale = _enum_value(mujoco.mjtFontScale, "mjFONTSCALE_100")
maxgeom = int(model.ngeom + len(path_points_visible) + 8)
scene = mujoco.MjvScene(model, maxgeom=maxgeom)
option = mujoco.MjvOption()
perturb = mujoco.MjvPerturb()
camera = mujoco.MjvCamera()
mujoco.mjv_defaultOption(option)
mujoco.mjv_defaultPerturb(perturb)
mujoco.mjv_defaultCamera(camera)
camera.type = camera_type
camera.fixedcamid = int(args.camera_id)
if hasattr(camera, "trackbodyid"):
camera.trackbodyid = -1
context = mujoco.MjrContext(model, font_scale)
viewport = mujoco.MjrRect(0, 0, args.width, args.height)
mujoco.mjv_updateScene(model, data, option, perturb, camera, catmask, scene)
for idx, path_point in enumerate(path_points_visible):
path_rgba[3] = 0.10 + 0.70 * (idx / max(len(path_points_visible) - 1, 1))
_append_sphere(scene, mujoco, path_point, 0.03, path_rgba)
rgb = np.empty((args.height, args.width, 3), dtype=np.uint8)
depth = np.empty((args.height, args.width), dtype=np.float32)
mujoco.mjr_render(viewport, scene, context)
mujoco.mjr_readPixels(rgb, depth, viewport, context)
imageio.imwrite(args.output_path, np.flipud(rgb))
context.free()
finally:
ctx.free()
bundle.env.close()
if __name__ == "__main__":
main()

10
scripts/poster_visualisations/render_static_path_image.sh Executable file
View file

@ -0,0 +1,10 @@
#!/usr/bin/env bash
# Static path image + optional ghost render
path=$1 # path to .flax model with metadata.yaml alongside it
uv run scripts/poster_visualisations/render_static_path_image.py \
"$path" \
--output-path vids/poster/5arms/centralized/path.png \
--ghost-overlay

3
scripts/simulate.py
View file

@ -107,6 +107,9 @@ def main(dict_cfg: DictConfig) -> None:
action_mask=action_mask,
output_path=output_path,
camera_id=sim_cfg.camera_id,
width=sim_cfg.video_width,
height=sim_cfg.video_height,
fps=sim_cfg.video_fps,
)
save_evaluation_metadata(

3
src/brittle_star_project/configs/config_simulation.py
View file

@ -26,6 +26,9 @@ class SimulationSettings:
video_output_path: Optional[str] = None
# Camera ID to use for video recording (1 is usually the close-up camera)
camera_id: int = 1
video_width: int = 640
video_height: int = 480
video_fps: int = 60
# Optional override for the sidecar metadata YAML file.
# If None, it defaults to the model_path with a `_metadata.yaml` suffix.

14
src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py
View file

@ -65,11 +65,21 @@ class BrittleStarJaxEnvWrapper:
def single_observation_space(self):
return self._env.observation_space
def reset(self, seed: int = 0):
def reset(self, seed: int = 0, target_position: tuple[float, float] | None = None):
self.logger.info(f"Resetting vectorized environment environments with seed {seed}")
self._action_rng, env_rng = jax.random.split(jax.random.PRNGKey(seed), 2)
env_rngs = jnp.array(jax.random.split(env_rng, self._num_envs))
state = self._vectorized_reset(rng=env_rngs)
# If a target_position is provided, pass it through to the underlying env.reset
if target_position is None:
state = jax.jit(jax.vmap(lambda rng: self._env.reset(rng=rng)))(env_rngs)
else:
tp = jnp.asarray(target_position, dtype=jnp.float32)
tp_batched = jnp.tile(tp[None, :], (self._num_envs, 1))
state = jax.jit(jax.vmap(lambda rng, t: self._env.reset(rng=rng, target_position=t)))(
env_rngs, tp_batched
)
return state
def sample_actions(self):

7
src/brittle_star_project/environment/env_wrapper.py
View file

@ -62,9 +62,12 @@ class BrittleStarEnv:
return jax.random.PRNGKey(seed)
def reset(self, *, seed: int = 0):
def reset(self, *, seed: int = 0, target_position: tuple[float, float, float] | None = None):
rng = self.make_rng(seed)
state = self._env.reset(rng=rng)
if target_position is not None:
state = self._env.reset(rng=rng, target_position=target_position)
else:
state = self._env.reset(rng=rng)
return state
def render(self, *, state: Any):

188
src/brittle_star_project/evaluation/video.py
View file

@ -25,6 +25,59 @@ def create_evaluation_dir(model_path: Path) -> Path:
return eval_dir
def _ensure_offscreen_size(model, width: int, height: int) -> None:
vis_global = getattr(getattr(model, "vis", None), "global_", None)
if vis_global is None:
return
vis_global.offwidth = int(max(width, vis_global.offwidth))
vis_global.offheight = int(max(height, vis_global.offheight))
def _apply_camera_overrides(
model,
*,
camera_fovy: dict[int, float] | None = None,
camera_xyz: tuple[
dict[int, float] | None,
dict[int, float] | None,
dict[int, float] | None,
] = (None, None, None),
) -> None:
if not camera_fovy and not (camera_xyz[0] or camera_xyz[1] or camera_xyz[2]):
return
ncam = int(getattr(model, "ncam", 0))
for cam_id, fovy in (camera_fovy or {}).items():
if cam_id < 0 or cam_id >= ncam:
raise ValueError(f"Camera id {cam_id} is out of range")
model.cam_fovy[cam_id] = float(fovy)
for cam_id, x in (camera_xyz[0] or {}).items():
if cam_id < 0 or cam_id >= ncam:
raise ValueError(f"Camera id {cam_id} is out of range")
model.cam_pos[cam_id][0] = float(x)
for cam_id, y in (camera_xyz[1] or {}).items():
if cam_id < 0 or cam_id >= ncam:
raise ValueError(f"Camera id {cam_id} is out of range")
model.cam_pos[cam_id][1] = float(y)
for cam_id, z in (camera_xyz[2] or {}).items():
if cam_id < 0 or cam_id >= ncam:
raise ValueError(f"Camera id {cam_id} is out of range")
model.cam_pos[cam_id][2] = float(z)
def hex_to_rgba(hex_color: str, alpha: float) -> np.ndarray:
color = hex_color.lstrip("#")
if len(color) != 6:
raise ValueError(f"Expected a 6-digit hex color, got {hex_color!r}")
red = int(color[0:2], 16) / 255.0
green = int(color[2:4], 16) / 255.0
blue = int(color[4:6], 16) / 255.0
return np.asarray([red, green, blue, float(alpha)], dtype=np.float32)
def save_evaluation_metadata(
eval_dir: Path,
*,
@ -66,6 +119,7 @@ def record_episode(
fps: int = 60,
width: int = 640,
height: int = 480,
target_xy: tuple[float, float] | None = None,
) -> EpisodeResult:
"""Run an episode headlessly and record a video using MuJoCo's Renderer and imageio.
@ -92,10 +146,12 @@ def record_episode(
"Please install the evaluation dependencies: `uv pip install .[evaluation]`"
) from e
state = env.reset(seed=seed)
state = env.reset(seed=seed, target_position=target_xy)
model = state.mj_model
data = state.mj_data
_ensure_offscreen_size(model, width, height)
renderer = mujoco.Renderer(model, width=width, height=height)
ep_return = 0.0
observations = _get_observations(state)
@ -147,3 +203,133 @@ def record_episode(
final_xy_dist=final_dist,
initial_target_distance=initial_dist,
)
def record_episode_multi_camera(
*,
env: BrittleStarEnv,
policy: ControlPolicy,
seed: int,
max_steps: int,
action_low: np.ndarray | None,
action_high: np.ndarray | None,
output_paths: dict[int, Path],
action_mask: np.ndarray | None = None,
camera_ids: list[int] | None = None,
camera_fovy: dict[int, float] | None = None,
camera_xyz: tuple[
dict[int, float] | None,
dict[int, float] | None,
dict[int, float] | None,
] = (None, None, None),
target_xy: tuple[float, float] | None = None,
robot_color: str | None = None,
fps: int = 60,
width: int = 640,
height: int = 480,
) -> EpisodeResult:
"""Run one episode and render multiple camera views to separate files."""
try:
import imageio
import mujoco
except ImportError as e:
raise ImportError(
"Video recording requires 'imageio' and 'mujoco'. "
"Please install the evaluation dependencies: `uv pip install .[evaluation]`"
) from e
if camera_ids is None:
camera_ids = list(output_paths.keys())
for cam_id in camera_ids:
if cam_id not in output_paths:
raise ValueError(f"Missing output path for camera {cam_id}")
output_paths = {cam_id: output_paths[cam_id] for cam_id in camera_ids}
for path in output_paths.values():
path.parent.mkdir(parents=True, exist_ok=True)
state = env.reset(seed=seed, target_position=(target_xy[0], target_xy[1], 0.0))
model = state.mj_model
data = state.mj_data
_apply_camera_overrides(model, camera_fovy=camera_fovy, camera_xyz=camera_xyz)
if robot_color is not None:
robot_body_id = mujoco.mj_name2id(
model, mujoco.mjtObj.mjOBJ_BODY, "BrittleStarMorphology/central_disk"
)
if robot_body_id < 0:
raise ValueError("Body 'BrittleStarMorphology/central_disk' not found in the model")
robot_rgba = hex_to_rgba(robot_color, 1.0)
body_parent = model.body_parentid
robot_body_ids = {int(robot_body_id)}
for body_id in range(1, int(model.nbody)):
current_body_id = int(body_id)
while current_body_id not in (-1, 0, int(robot_body_id)):
current_body_id = int(body_parent[current_body_id])
if current_body_id == int(robot_body_id):
robot_body_ids.add(body_id)
for geom_id in range(int(model.ngeom)):
if int(model.geom_bodyid[geom_id]) in robot_body_ids:
model.geom_rgba[geom_id][:] = robot_rgba
_ensure_offscreen_size(model, width, height)
renderer = mujoco.Renderer(model, width=width, height=height)
writers = {
cam_id: imageio.get_writer(str(path), fps=fps) for cam_id, path in output_paths.items()
}
ep_return = 0.0
observations = _get_observations(state)
prev_dist = _get_xy_distance_to_target(observations) if observations else None
initial_dist = prev_dist
reached_target = _target_reached(state=state)
steps = 0
try:
for _ in range(int(max_steps)):
for cam_id in camera_ids:
renderer.update_scene(data, camera=cam_id)
writers[cam_id].append_data(renderer.render())
obs_dict = observations or {}
action = policy.act(observations=obs_dict)
if action_mask is not None:
action = action[action_mask]
action = _maybe_clip_action(action, action_low, action_high)
state = env.step(state=state, action=action)
steps += 1
observations = _get_observations(state)
cur_dist = _get_xy_distance_to_target(observations) if observations else None
if prev_dist is not None and cur_dist is not None:
ep_return += prev_dist - cur_dist
prev_dist = cur_dist
reached_target = _target_reached(state=state)
if reached_target:
break
for cam_id in camera_ids:
renderer.update_scene(data, camera=cam_id)
writers[cam_id].append_data(renderer.render())
finally:
renderer.close()
for writer in writers.values():
writer.close()
final_dist = _get_xy_distance_to_target(observations) if observations else None
return EpisodeResult(
return_=ep_return,
length=steps,
reached_target=reached_target,
final_xy_dist=final_dist,
initial_target_distance=initial_dist,
)