We are going to define some simulation environments.

In each case, z is vertical, x is lateral, y is depth. We'll use right hand rule, so +z is up, +x is right, +y is forward (typically down into the screen).

Each vehicle works like this:
- implements the base_dynamics.py abstract class, with a custom class that defines how the state updates based on control inputs: step(self, state, action), returning next_state, debug_dictionary. It should also implement state_variable_descriptions() returning a list of strings describing each state variable, and state_variable_symbols() returning a list of strings with single character symbols for each state variable (unicode). The indices of the states to visualize as manifolds should be returned by state_indices_manifold(). You should also have descriptions and symbols for the control inputs as well. dt and physical_parameters should be constructor parameters, define what is in physical_parameters in the docstring (what it is, units if possible, and what happens if you increase or decrease it)
- implements a graphics class that defines how to visualize the environment in 3D using pygfx
- is wrapped in a SeamstressEnvironment class that handles resetting, stepping, and rendering the environment. The environment class defines the state_limits and control_limits

----------------------------------------------------------------

seamstress/src/seamstress/dynamics/rotorcraft.py

The rotorcraft dynamics will be simple and constrained to the surface of a 3d hemispherical like object, simulating a drone gathering images of an object.

constructor:
- imaging_radius: the radius of the hemisphere the rotorcraft is flying around

state: 
- angle_around_vertical_deg
- angle_from_vertical_deg

control inputs:
- angular_velocity_around_vertical
- angular_velocity_from_vertical

we should incorporate mass, and inertia, and have a simple drag model based on the vehicles velocity.

When returning the new state, return a debug_dictionary with the following entries:
- F_tangential (as a 3D vector in world frame)
- F_drag   
- F_total

seamstress/src/seamstress/graphics/rotorcraft.py

body_static should return a pygfx object that contains the following primitives defining the vehicle:
- four flat cylinders each of radius extent_horizontal/6 and thickness extent_vertical, representing the rotors, placed at the four cardinal directions around the vehicle, offset from the center by extent_horizontal/2 in x and y.
- a flat cuboid of shape (extent_horizontal, extent_horizontal, extent_vertical/6) representing the body of the rotorcraft, centered at the origin.

body_dynamic should return a pygfx object that contains a cylinder representing the camera gimbal, of radius extent_horizontal/8, and height extent_horizontal/2. One end should be centered at the origin, and it should extend out along the -z axis. This object will be rotated during simulation to always point at the center of the hemisphere.

scene_static should return a pygfx object that contains:
- a red cuboid of shape (length, width, height) representing the object being imaged, centered at the origin.
- a grey plane representing the ground plane, it should be imaging radius*4 in x and y, and offset vertically -cuboid_height/2 so that the cuboid is sitting on top of it.

views should return a dictionary of CameraView objects defining the following camera views for rendering the environment:
- "ISOMETRIC": An isometric orthographic view showing the rotorcraft from a 45 degree angle up from the horizontal, and 45 degree angle around the vertical. It should be at imaging radius * 1.5 distance from the origin. z should be up.
- "EGO": An ego-centric perspective view simulating the camera mounted on the rotorcraft, always looking at the center of the hemisphere. The camera should be positioned at the rotorcraft position, looking at the origin, with +z up.
- "BIRDSEYE": A top-down orthographic view showing the rotorcraft from above, looking straight down along the -z axis. The camera should be positioned at the origin, with +y and +x defining the up and right directions respectively in the rendered image.

----------------------------------------------------------------

seamstress/src/seamstress/dynamics/glider.py

The glider flies from left to right, with a single control surface (elevator) on the wings on the aircraft. The elevator can be deflected up or down to control the pitch of the glider.

constructor:
- None 

state:
- x_position
- z_position
- pitch_angle_deg (degrees, 0 is horizontal, positive is nose up)

control inputs:
- aileron_deflection_deg (degrees, positive is trailing edge down)

Implement gravitational effects, and the aircraft produces a lift and drag force based on its speed and pitch, as well as a pitching torque based on the aileron deflection. If the pitch angle of the glider exceeds a certain threshold, enter a stall condition where lift is drastically reduced.

When returning the new state, return a debug_dictionary with the following entries:
- F_lift (as a 2D vector in world frame)
- F_drag
- F_gravity
- F_total
- torque 

seamstress/src/seamstress/graphics/glider.py

body_static should return a pygfx object that contains the following primitives defining the vehicle:
- a cuboid representing the main body of the glider, with shape (length, width/10, height/2)
- a triangle tail fin at the rear of the glider of height height/2 and base height/2
- a cuboid representing the wings, centered at the origin with shape (length/4, width, height/10)

body_dynamic should return a pygfx object that contains the following primitives defining the control surface:
- an aileron on the trailing edge of the wings, represented as a thin cuboid of shape (length/8, width, height/10), hinged at the back of the wings.

scene_static should return a pygfx object that contains:
- a grey ground plane placed according to the environment's state limits: a thin strip at state_limit.z_min extending from state_limit.x_min to state_limit.x_max, and glider_width in y.

There should be a ground plane placed according to the environment's terminator.state_limits. There should also be walls placed according the environment's terminator.state_limits, extending from state_limits.x_min to state_limits.x_max, and from z=0. Do not place walls if the state_limit is np.inf or -np.inf for that axis.

The camera views should be:
- "ISOMETRIC": An isometric orthographic view showing the glider from a 45 degree angle about the vertical, and 45 degree angle above the horizontal. z should be up. moves with the glider.
- "SIDE": An ego-centric isometric orthographic view centered on the glider (i.e. always offset the same amount from the glider), looking forward along the y-axis, with +z up, fairly zoomed out.

----------------------------------------------------------------

For all environments, implement reasonable default parameters. Use a constructor-defined time step for dynamics updates (e.g., 0.1s). Implement an Integrator class that implements Euler and RK4 for updating state based on velocities and accelerations. Ignore external disturbances like wind for now. IMPORTANTLY: all dynamics functions should be implemented with jax so we can have autograd and JIT for speed.

Implement base_dynamics.py which is an ABC base class that all the above inherit from. It should define abstract methods:
- step(self, state, action): takes a state and action input, returns the next state after one time step
- __init__(self, parameters={}): initializes the environment with given parameters or defaults

Implement base_graphics.py which is an ABC base class that all the above inherit from. It should define abstract methods:
- body_static(self): returns a pygfx object representing the static parts of the vehicle
- body_dynamic(self): returns a pygfx object representing the dynamic parts of the vehicle that move
- scene_static(self): returns a pygfx object representing the static parts of the environment (e.g., ground plane, obstacles)
- views(self): returns a dictionary of CameraView objects defining different camera views for rendering the environment

Also implement the view.py file that defines the camera type, fov, position, orientation, image size, rendering utilities, etc.

Implement seamstress/src/seamstress/environments/environment.py that defines the SeamstressEnvironment class which wraps around a dynamics and graphics class. It should implement:
- reset(self): resets the environment to a default state
- step(self, action): applies the action to the environment, updates the state using the dynamics, and returns the new state, reward, done flag, and info dictionary
- render(self, view_name): renders the environment from the specified camera view and returns the rendered image as a numpy array
- __init__(self, dynamics: BaseDynamics, graphics: BaseGraphics, reset_fn, state_limits, control_limits): initializes the environment with given dynamics and graphics classes, as well as state and control limits

DO NOT IMPLEMENT ANY OTHER CLASSES OR FILES IN THE TREE BELOW, but do feel free to suggest utility files to put in utils if you are repeating code, or for neatness/organisation.

For example, you might like to implment this in utils/conversions.py:
<<<<CONVERSIONS START<<<<
# put this wherever you like, or duplicate in each dynamics file
import jax.numpy as jnp

DEG2RAD = jnp.pi / 180.0
RAD2DEG = 180.0 / jnp.pi
>>>>CONVERSIONS END<<<<

----------------------------------------------------------------

To test all of these, implement a file called seamstress/src/tests/test_environments.py that creates an instance of each environment, resets it to a default state, applies a series of random valid actions until the episode ends (according to the terminator), and renders the environment at each step, saving the rendered images to disk.

Use the tools in my logging tools (at the end of this prompt src/seamstress/utils/custom_logging.py) and the video tools (at end of this prompt src/seamstress/utils/video.py) to log the rendered images as videos to W&B. The logging folder should look like this:
seamstress/
  logs/
    run_xyz/
        <environment>_images/
            frame_00001.png
            frame_00002.png
            ...
        <environment>_video.mp4

With the videos also logged to wandb.

I know that in practice, we will need the first derivatives of the three suggested state variables for each environment, so it's ok to have more state variables than those listed above if needed for proper dynamics simulation. But I will only ever visualize as manifolds the three suggested state variables for each environment.

My main is structured like so:
>>>>MAIN START<<<<
#!/usr/bin/env python3
from omegaconf import OmegaConf, DictConfig
import hydra
import inspect
import argparse
import os
import shutil

@hydra.main(config_path="conf", config_name="config", version_base=None)
def main(args: DictConfig):
    # Log out the configuration
    yaml_str = OmegaConf.to_yaml(args)
    print(f"[{inspect.stack()[0][3]}] configuration:\n{yaml_str}")
    
if __name__ == "__main__":
    main()
>>>>MAIN END<<<<

And you should structure the test code like this so I can call it from the command line. You are allowed to modify the config.yaml file to add anything you need. In general do not hardcode, always use config wherever possible.

ubuntu@192-222-59-201:~/user_irw/seamstress$ tree src/seamstress/
src/seamstress/
├── __init__.py
├── __pycache__
│   └── __init__.cpython-311.pyc
├── conf
│   ├── __init__.py
│   └── config.yaml
├── control
│   ├── mppi.py
│   └── random_sampler.py
├── dynamics
│   ├── __pycache__
│   │   ├── base_dynamics.cpython-311.pyc
│   │   └── rotorcraft.cpython-311.pyc
│   ├── glider.py
│   ├── rotorcraft.py
│   └── utils
│       ├── __pycache__
│       │   └── integrator.cpython-311.pyc
│       ├── base_dynamics.py
│       └── integrator.py
├── environments
│   ├── __pycache__
│   │   ├── environment.cpython-311.pyc
│   │   └── terminator.cpython-311.pyc
│   └── environment.py
├── graphics
│   ├── __pycache__
│   │   └── video.cpython-311.pyc
│   ├── glider.py
│   ├── rotorcraft.py
│   └── utils
│       ├── base_primitive.py
│       └── view.py
├── main.py
├── manifolds
│   ├── curvature.py
│   ├── manifold.py
│   └── tangent_bundle.py
├── tests
│   ├── __pycache__
│   │   └── test_environments.cpython-311.pyc
│   └── test_environments.py
└── utils
    ├── __pycache__
    │   └── custom_logging.cpython-311.pyc
    ├── cacher.py
    ├── custom_logging.py
    ├── custom_random.py
    ├── video.py
    └── wandb_cleanup.py

>>>>LOGGING TOOLS START<<<<
import os 
import time 
import datetime
import pathlib
import tempfile
import shutil
import wandb
import numpy as np
import matplotlib.pyplot as plt

def get_repo_root_dir():
    # We're in repo/src/utils/logging.py
    return os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

def get_timestamp():
    return datetime.datetime.now().strftime("%Y_%m_%d_%H_%M_%S")

def make_log_dir():
    log_dir = os.path.join(get_repo_root_dir(), "logs")

    # Make a timestamped subdirectory
    foldername = f"run_{get_timestamp()}"
    full_path = os.path.join(log_dir, foldername)
    os.makedirs(full_path, exist_ok=True)
    # Return as a pathlib.Path object
    return pathlib.Path(full_path)

def get_saved_dir():
    return os.path.join(get_repo_root_dir(), "saved")

def get_cache_dir():
    return os.path.join(get_saved_dir(), "cache")

def get_weights_dir():
    return os.path.join(get_saved_dir(), "weights")

def log_figure_to_wandb(figure, key: str):
    """
    Save a Matplotlib figure to a temporary PNG (dpi=600), log it to W&B, then close the figure.
    """
    # Create a temp file path
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
    tmp_path = tmp.name
    tmp.close()

    try:
        # Save figure at high resolution
        figure.savefig(tmp_path, dpi=300, bbox_inches="tight")
        # Log to wandb with global step
        wandb.log({key: wandb.Image(tmp_path)})
    finally:
        # Close and clean up
        plt.close(figure)
        try:
            os.remove(tmp_path)
        except OSError:
            pass

def log_video_to_wandb(clip, key: str):
    """
    Save a MoviePy clip to a temporary MP4, log it to W&B, then clean up.
    Uses the clip's own fps/size.
    """
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
    tmp_path = tmp.name
    tmp.close()

    try:
        clip.write_videofile(
            tmp_path,
            codec="libx264",
            audio=False,
            logger=None
        )
        wandb.log({key: wandb.Video(tmp_path, format="mp4")})
    finally:
        try:
            clip.close()   # make sure resources (ffmpeg, readers) are freed
        except Exception:
            pass
        try:
            os.remove(tmp_path)
        except OSError:
            pass
>>>>LOGGING TOOLS END<<<<

>>>>VIDEO TOOLS START<<<<
# utils/image_folder_to_video.py

from __future__ import annotations

import os
from pathlib import Path
from typing import List

from moviepy.video.io.ImageSequenceClip import ImageSequenceClip

IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff"}


def _collect_images(folder: Path) -> List[Path]:
    """
    Collect frame files from a directory, sorted numerically when possible.
    """
    if not folder.is_dir():
        raise FileNotFoundError(f"Image folder does not exist: {folder}")

    files = [
        p for p in folder.iterdir()
        if p.is_file() and p.suffix.lower() in IMAGE_EXTENSIONS
    ]
    if not files:
        raise FileNotFoundError(f"No image files found in: {folder}")

    def sort_key(path: Path):
        stem = path.stem
        try:
            return (0, int(stem))  # numeric sort
        except ValueError:
            return (1, stem)       # fallback lexicographic

    files.sort(key=sort_key)
    return files


def image_folder_to_video(
    filepath_image_folder: str | os.PathLike,
    filepath_output_video: str | os.PathLike,
    *,
    fps: int = 24,
) -> Path:
    """
    Generic converter: turn a folder of images into an MP4 video.

    Parameters
    ----------
    filepath_image_folder : str | Path
        Directory containing image frames.
    filepath_output_video : str | Path
        Desired output video path. If missing '.mp4', a warning is printed.
    fps : int, default=24
        Frames per second.

    Returns
    -------
    Path
        Path to the written MP4 file.
    """
    folder = Path(filepath_image_folder)
    output = Path(filepath_output_video)

    # Warn if user forgot extension
    if output.suffix.lower() != ".mp4":
        print(f"[warning] Output filename does not end with '.mp4': {output}")

    image_files = _collect_images(folder)
    clip = ImageSequenceClip([str(p) for p in image_files], fps=fps)

    clip.write_videofile(
        str(output),
        fps=fps,
        codec="libx264",
        audio=False,
    )
    clip.close()
    return output
>>>>VIDEO TOOLS END<<<<