Co-Design
Gym

import codesign_gym as cg

# one entry point; returns the env and its base design
env, theta0 = cg.make_vec("BallCatcher-Pitch", num_envs=8)

# your outer loop proposes a design ...
theta = optimizer.ask(env.embodiment_space, theta0)
obs, rejected, info = env.reset(embodiment=theta)

# ... and your inner loop acts in it
for t in range(horizon):
    obs, reward, term, trunc, info = env.step(policy(obs, theta))

# reset returns a per-env rejection flag instead of raising
obs, rejected, info = env.reset(embodiment=thetas)
if rejected.any():
    # structurally impossible designs: out of bounds,
    # wrong shape, overlapping cells. Score them and re-propose;
    # stepping a rejected env is an error.
    optimizer.tell(thetas[rejected], penalty)

# check a design without touching the simulator state
ok = env.validate_embodiment(theta)

# Racing and the Shape-shifting Hand accept a design at every step
env, theta = cg.make_vec("RotateInHand", num_envs=64)
obs, _, _ = env.reset(embodiment=theta)
for t in range(horizon):
    theta = theta + 0.01 * design_policy(obs)   # reshape the pads
    obs, r, term, trunc, info = env.step(action, embodiment=theta)

# an eval env ignores step-time designs (with a warning),
# so a fixed design cannot be reshaped mid-episode
eval_env, _ = cg.make_vec("RotateInHand", 8, is_eval_env=True)

# envs/my_env/registration.py
def build_vec(num_envs, *, is_eval_env, **cfg):
    env = MyNativeVectorEnv(num_envs, **cfg)  # vmap, shards, warp...
    env.base_embodiment = default_design()
    return env

# codesign_gym/registry.py: one row, imported only when used
REGISTRY["MyEnv-v0"] = ("envs.my_env.registration:build_vec", {})