f594cafead
Provide the features to read in Environment (Morphology, arena, ...) config files in the JSON format. The example JSON file contains the config for a brittle star with 2 arms
95 lines
2.7 KiB
Python
95 lines
2.7 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
from brittle_star_project import (
|
|
Backend,
|
|
BrittleStarEnv,
|
|
BrittleStarEnvFactory,
|
|
SimulationConfig,
|
|
simulate_policy,
|
|
)
|
|
from brittle_star_project.environment import from_json
|
|
from brittle_star_project.rl import RLModel # imports concrete models via rl.__init__
|
|
from brittle_star_project.rl.base import get_rl_model_registry
|
|
|
|
MODEL_BY_NAME = get_rl_model_registry()
|
|
MODEL_OPTIONS = sorted(MODEL_BY_NAME)
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
p = argparse.ArgumentParser(description="Simulate a trained policy in the MuJoCo viewer.")
|
|
p.add_argument(
|
|
"--model",
|
|
type=str,
|
|
default=None,
|
|
help="Path to a saved model artifact. If omitted, a model is created from --model-type.",
|
|
)
|
|
p.add_argument(
|
|
"--model-type",
|
|
choices=MODEL_OPTIONS,
|
|
default="random",
|
|
help="Which model class to instantiate when --model is omitted.",
|
|
)
|
|
p.add_argument(
|
|
"--backend",
|
|
choices=[b for b in Backend],
|
|
default=Backend.MJX,
|
|
)
|
|
p.add_argument("--seed", type=int, default=None)
|
|
return p.parse_args()
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
|
|
morphology_cfg, arena_cfg, env_cfg = from_json("../configs/test.json")
|
|
|
|
# ======= ENVIRONMENT SETUP =======
|
|
|
|
backend = args.backend
|
|
|
|
factory = BrittleStarEnvFactory()
|
|
raw_env = factory.create_environment(backend, morphology_cfg, arena_cfg, env_cfg)
|
|
env = BrittleStarEnv(raw_env, backend=backend, config=env_cfg)
|
|
|
|
seed_for_env = int(args.seed) if args.seed is not None else 0
|
|
state = env.reset(seed=seed_for_env)
|
|
|
|
# ======= MODEL SETUP =======
|
|
|
|
# Extract the number of actuators (nu) from the environment's model, so we can pass it to the
|
|
# policy/model.
|
|
nu = int(state.mj_model.nu)
|
|
|
|
if args.model is not None:
|
|
model_path = Path(args.model)
|
|
policy = RLModel.load(model_path)
|
|
if hasattr(policy, "nu"):
|
|
policy.nu = nu
|
|
else:
|
|
model_cls = MODEL_BY_NAME[str(args.model_type)]
|
|
policy = model_cls(seed=seed_for_env)
|
|
if hasattr(policy, "nu"):
|
|
policy.nu = nu
|
|
|
|
# If the policy/model has a `seed` attribute, use the provided seed (or default) to reset it.
|
|
default_seed = int(getattr(policy, "seed", seed_for_env))
|
|
if args.seed is not None and hasattr(policy, "reset"):
|
|
policy.reset(int(args.seed))
|
|
|
|
# ======= SIMULATION =======
|
|
|
|
rollout_cfg = SimulationConfig(
|
|
realtime=True,
|
|
seed=int(args.seed) if args.seed is not None else default_seed,
|
|
)
|
|
|
|
simulate_policy(policy, rollout_cfg, state)
|
|
|
|
env.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|