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 environment works like this:
- implements the base_dynamics.py abstract class, with a custom class that defines how the state updates based on control inputs
- 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.

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

MANIPULATOR

(torus/cylinder)

DYNAMICS
A 3d manipulator arm without an end effector. The base can be rotated about z, the first joint can move up and down (pitch), and the second joint can also move up and down (pitch). The arm segments have fixed lengths.
state_variables = {base_about_z_angle, joint_1_angle, joint_2_angle}, when all are zero, the arm is fully extended along the +z axis (straight up). the first state variable rotates clockwise about z relative to +y, the other two are pitch angles relative to the previous segment, all expressed in degrees.
action_inputs = {base_about_z_w, joint_1_w, joint_2_w} (angular velocities)
parameters
- base_height
- base_radius
- link_1_length
- link_1_radius
- link_2_length
- link_2_radius

Nothing too fancy here, just move the joints according to the input angular velocities. State limits will be provided by the terminator to prevent impossible configurations.

When returning the new state, return a debug_dictionary with the following entries:
- end_effector_position (x, y, z) in world frame
- joint_1_position (x, y, z) in world frame
- joint_2_position (x, y, z) in world frame

ENVIRONMENT 
The 3D visualization of the manipulator should consist of:
- A cylindrical base at the origin with height base_height and radius base_radius
- A cylindrical first link extending from the top of the base, with length link_1_length and radius link_1_radius
- A cylindrical second link extending from the end of the first link, with length link_2_length and radius link_2_radius
This object should be constructed from primitive shapes in pygfx, and rotated according to the state variables.

The manipulator should be in a box cube of side length equal to twice the sum of the link lengths, centered at the origin. There should be another box representing a table beneath it, with height half the length of the outer cube, and x and y extents equal to half the outer cube's extents. The base of the manipulator should be on top of the table (i.e. the origin).

The camera views should be:
- "ISOMETRIC": An isometric orthographic view showing the manipulator from a 45 degree angle above and to the side.
- "EGO": An ego-centric isometric orthographic view looking down the end of link2, mimicing a camera mounted at the end of the manipulator arm.
- "SCENE": A front perspective view centered on the origin, looking directly along the +y axis, with +z up and the whole scene in view.

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

ROCKET

(cone)

DYNAMICS
A single stage rocket that can thrust upwards and gimbal its engine to control pitch, and is trying to land upright. Should NOT implement fuel mass (this would require a fourth state) and some rocket equations.
state_variables = {x, z, pitch}, where pitch is the angle clockwise from the vertical +z expressed in degrees
action_inputs = {thrust, gimbal_angle}, where gimbal_angle is the angle the thrust vector makes with the vertical, clockwise from +z, expressed in degrees.
parameters
- extents (x,z)
- mass (fuel and structure) (you should be able to run out of fuel)
- various aerodynamic coefficients
- rotation matrix type stuff

Implement gravitational effects, and the rocket produces a thrust force based on the thrust input and a pitching torque based on the gimbal angle. Also implement a simple drag model based on the rocket's velocity and pitch.

When returning the new state, return a debug_dictionary with the following entries:
- F_thrust (should be in the direction defined by the gimbal angle), (as a 2D vector in world frame)
- F_drag
- F_gravity (straight down)
- F_total (the vector sum of all forces)
- torque (should be based on the gimbal angle and the distance from the center of mass)

ENVIRONMENT 
The 3D visualization of the rocket should be a cuboid body of shape (extents_x, extents_x, extents_z). This object should be constructed from primitive shapes in pygfx, and rotated according to the pitch state variable. Make the top triangular (taking up a quarter of the height) to represent a nose cone.

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 rocket from a 45 degree angle above and to the side.
- "EGO": An ego-centric isometric orthographic view from the perspective of the rocket facing downwards along it's local -z axis, i.e. a landing camera view.
- "SCENE": A front perspective view centered on the origin, looking directly along the +y axis, with +z up and the whole scene in view.

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

SATELLITE

DYNAMICS
A cube satellite with 8 thrusters, one mounted at each corner. No external forces or torques (ignore gravity for now). The +z face is the front of the satellite (the 'camera' points along +z, and the solar panels face along +z). The thrusters are mounted so that their axis of thrust is along the x axis. The ones on the +x side point +x, the ones on the -x side point -x. Firing a thruster produces a force along its axis of thrust at the location of the thruster, producing a torque about the center of mass.

We are concerned with it's attitude (not it's position). The satellite starts with the camera and solar panels facing +z.
state_variables = {roll, pitch, yaw}, all expressed in degrees.
action_inputs = {thruster_1, thruster_2, thruster_3, thruster_4, thruster_5, thruster_6, thruster_7, thruster_8}, each thruster input is a binary on/off command.
parameters
- extents (length of one edge of the cube)
- mass
- moment_of_inertia (can be a scalar, assume uniform mass distribution)
- thruster_force (the force each thruster produces when activated)

Do not implement gravitational effects or drag. All we are concerned with is the attitude control of the satellite using the thrusters. If a thruster fires, compute the torque it produces about the center of mass and update the angular velocities accordingly. Do this for all thrusters that are activated in the action input.

When returning the new state, return a debug_dictionary with the following entries:
- T (total torque about the center of mass, as a 3D vector in world frame)
- F_thruster_1 (as a 3D vector in world frame)
- F_thruster_2 (as a 3D vector in world frame)
- F_thruster_3 (as a 3D vector in world frame)
- F_thruster_4 (as a 3D vector in world frame)
- F_thruster_5 (as a 3D vector in world frame)
- F_thruster_6 (as a 3D vector in world frame)
- F_thruster_7 (as a 3D vector in world frame)
- F_thruster_8 (as a 3D vector in world frame)

GRAPHICS
The 3D visualization of the satellite should be a cube of shape (extents, extents, extents), with two rectangular solar panels extending from the +y and -y faces (not the faces with the thrusters). Each solar panel should be constructed of two pieces, a cylindrical arm of radius extents/10 and length extents/2 extending from the center of the face (colored dark gray), and a rectangular panel half the width as the satellite, 4*extents in length, and with thickess extents/20, attached to the end of the arm (colored dark blue). This object should be constructed from primitive shapes in pygfx, and rotated according to the roll, pitch, and yaw state variables. There should be a flat disk (dark gray) on the +z face representing the camera, with radius extents/2 and thickness extents/8.

The satellite (light gray/silver) should have a small cylinder representing each thruster at each corner, each thruster having a radius of extents/8 and a height of extents/4, positioned so that they are flush with the corners of the cube and extend outwards.

There should not be any walls, but the scene should be a radius 100 black sphere, so it looks like space.

The camera views should be:
- "ISOMETRIC": An isometric orthographic view showing the satellite from a 45 degree angle above and to the side.
- "EGO": An ego-centric isometric perspective view centered on the satellite (i.e. always offset the same amount from the satellite), looking forward along the camera direction, but mounted slightly above the +z axis so we can see the satellite body. 
- "SCENE": A front perspective view centered on the origin, looking directly along the +y axis, with +z up and the whole scene in view.

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

For all environments, implement reasonable default parameters. Use a fixed 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. Don't worry about rewards or goals, just focus on the dynamics, we will later (not now) define termination conditions and make these into environments. 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

We will use these in a SeamstressEnvironment class that implements the following methods:
- reset(self, state): resets the environment to the given state
- step(self, action): applies the action to the environment, updates the state using the underlying dynamics model, and returns the new state
- get_state(self): returns the current state of the environment
- render(self, view): renders a simple 3D visualization using pygfx, using a specified camera view. View is gonna be it's own class that defines camera position, orientation, ortho/persp, perspective_fov, orthographic_scale, image_width, image_height, etc.
- __init__(self, dynamics_model: BaseDynamics, terminator, control_limits, views=[]): initializes the environment with a specific dynamics model, termination conditions, and a list of camera views for rendering. Terminator is it's own class that defines termination conditions, initially, we'll just implement num_step_limit and state_limits (upper and lower bounds for each state variable).

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

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.

Do not attempt to implement anything in utils/, manifolds/, control/. Otherwise you can change things up, what I have here is just a guide. 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<<<<

Importantly, you must NOT edit the number of state variables - always 3 state variables.

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.

This is what the directory looks like:
ubuntu@192-222-59-201:~/user_irw/seamstress$ tree
.
├── README.md
├── prompt.txt
├── pyproject.toml
├── src
│   ├── seamstress
│   │   ├── __init__.py
│   │   ├── conf
│   │   │   ├── __init__.py
│   │   │   └── config.yaml
│   │   ├── control
│   │   │   ├── mppi.py
│   │   │   └── random_sampler.py
│   │   ├── dynamics
│   │   │   ├── base_dynamics.py
│   │   │   ├── glider.py
│   │   │   ├── manipulator.py
│   │   │   ├── rocket.py
│   │   │   ├── rotorcraft.py
│   │   │   ├── rover.py
│   │   │   ├── satellite.py
│   │   │   ├── submarine.py
│   │   │   └── utils
│   │   │       └── integrator.py
│   │   ├── environments
│   │   │   ├── environment.py
│   │   │   └── terminator.py
│   │   ├── main.py
│   │   ├── manifolds
│   │   │   ├── curvature.py
│   │   │   ├── manifold.py
│   │   │   └── tangent_bundle.py
│   │   ├── rendering
│   │   │   ├── primitives.py
│   │   │   ├── reference_frames.py
│   │   │   └── video.py
│   │   ├── tests
│   │   └── utils
│   │       ├── cacher.py
│   │       ├── custom_logging.py
│   │       ├── custom_random.py
│   │       └── wandb_cleanup.py
│   └── template.egg-info
│       ├── PKG-INFO
│       ├── SOURCES.txt
│       ├── dependency_links.txt
│       ├── requires.txt
│       └── top_level.txt
├── torch_versions.txt
└── uv.lock

>>>>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<<<<