Merge branch 'dev' into feat/wandb-logging
This commit is contained in:
commit
512272d6ab
SSH key fingerprint:
SHA256:u/h/LVoqKF1Iz02uOyxe6hcjmoZASCGV2HM0TG9ZMoU
41 changed files with 1517 additions and 792 deletions
|
|
@ -2,6 +2,7 @@ from .environment.env_types import Backend, Task
|
|||
from .environment.env_config import ArenaConfig, EnvConfig, MorphologyConfig
|
||||
from .environment.factory import BrittleStarEnvFactory
|
||||
from .environment.env_wrapper import BrittleStarEnv
|
||||
from .render import simulate_policy, SimulationConfig, ControlPolicy
|
||||
|
||||
__all__ = [
|
||||
"ArenaConfig",
|
||||
|
|
@ -11,4 +12,7 @@ __all__ = [
|
|||
"EnvConfig",
|
||||
"MorphologyConfig",
|
||||
"Task",
|
||||
"simulate_policy",
|
||||
"SimulationConfig",
|
||||
"ControlPolicy",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -1,15 +1,30 @@
|
|||
from dataclasses import dataclass
|
||||
|
||||
import jax
|
||||
|
||||
|
||||
@jax.tree_util.register_dataclass
|
||||
@dataclass
|
||||
class PPOArgs:
|
||||
"""
|
||||
source: https://github.com/vwxyzjn/cleanrl/blob/master/cleanrl/ppo_atari_envpool_xla_jax_scan.py
|
||||
"""
|
||||
|
||||
# path to environment config file, if None, use default config
|
||||
env_config_path: str | None = None
|
||||
|
||||
# path to hyperparameter config file (yaml), if None, use default config
|
||||
hyperparameter_config_path: str | None = None
|
||||
|
||||
# the name of this experiment
|
||||
exp_name: str = "brittle_star_ppo"
|
||||
|
||||
# the directory to save the experiment results
|
||||
run_dir: str | None = None
|
||||
|
||||
# how often to save checkpoints (0 to disable)
|
||||
checkpoint_frequency: int = 0
|
||||
|
||||
# seed of the experiment
|
||||
seed: int = 1
|
||||
|
||||
|
|
@ -44,8 +59,6 @@ class PPOArgs:
|
|||
hf_entity: str = ""
|
||||
|
||||
# ==== Algorithm specific dataclasses ====
|
||||
# the id of the environment
|
||||
env_id: str = "" # todo
|
||||
|
||||
# total timesteps of the experiments
|
||||
total_timesteps: int = 10000000
|
||||
|
|
@ -54,7 +67,7 @@ class PPOArgs:
|
|||
learning_rate: float = 2.5e-4
|
||||
|
||||
# the number of parallel game environments
|
||||
num_envs: int = 16
|
||||
num_envs: int = 100
|
||||
|
||||
# the number of steps to run in each environment per policy rollout
|
||||
num_steps: int = 128
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from brittle_star_project import (
|
|||
ArenaConfig,
|
||||
Backend,
|
||||
)
|
||||
from brittle_star_project.environment import from_file
|
||||
|
||||
|
||||
class BrittleStarJaxEnvWrapper:
|
||||
|
|
@ -84,3 +85,21 @@ class BrittleStarJaxEnvWrapper:
|
|||
return BrittleStarJaxEnvWrapper(
|
||||
morphology, arena, env_config, num_envs=num_envs, backend=backend
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def from_config(
|
||||
config_path: str, num_envs: int, backend: Backend = Backend.MJX
|
||||
) -> "BrittleStarJaxEnvWrapper":
|
||||
morphology_cfg, arena_cfg, env_cfg = from_file(config_path)
|
||||
return BrittleStarJaxEnvWrapper(
|
||||
morphology_cfg, arena_cfg, env_cfg, num_envs=num_envs, backend=backend
|
||||
)
|
||||
|
||||
def __str__(self):
|
||||
morphology_str = str(self._morphology)
|
||||
arena_str = str(self._arena)
|
||||
env_config_str = str(self._env_config)
|
||||
return (
|
||||
f"BrittleStarJaxEnvWrapper(backend={self._backend}, num_envs={self._num_envs}, "
|
||||
+ f"morphology={morphology_str}, arena={arena_str}, env_config={env_config_str})"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
from .env_config import ArenaConfig, EnvConfig, MorphologyConfig
|
||||
from .env_config import ArenaConfig, EnvConfig, MorphologyConfig, from_file
|
||||
from .env_types import Backend, Task
|
||||
from .env_wrapper import BrittleStarEnv, StepResult
|
||||
from .factory import BrittleStarEnvFactory
|
||||
|
|
@ -12,4 +12,5 @@ __all__ = [
|
|||
"BrittleStarEnv",
|
||||
"StepResult",
|
||||
"BrittleStarEnvFactory",
|
||||
"from_file",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
import json
|
||||
|
||||
from .env_types import Task
|
||||
|
||||
|
|
@ -48,6 +49,18 @@ class EnvConfig:
|
|||
# Per docs in upstream env config: integer factors of 200.
|
||||
light_perlin_noise_scale: int = 0
|
||||
|
||||
@staticmethod
|
||||
def from_json(path: str) -> EnvConfig:
|
||||
pass
|
||||
|
||||
def from_file(path: str) -> tuple[MorphologyConfig, ArenaConfig, EnvConfig]:
|
||||
"""Load configurations from a JSON or YAML file."""
|
||||
with open(path, "r") as f:
|
||||
if path.endswith(".yaml") or path.endswith(".yml"):
|
||||
import yaml
|
||||
|
||||
config_dict = yaml.safe_load(f)
|
||||
else:
|
||||
config_dict = json.load(f)
|
||||
|
||||
morphology = MorphologyConfig(**config_dict.get("morphology", {}))
|
||||
arena = ArenaConfig(**config_dict.get("arena", {}))
|
||||
env = EnvConfig(**config_dict.get("env", {}))
|
||||
return morphology, arena, env
|
||||
|
|
|
|||
|
|
@ -0,0 +1,3 @@
|
|||
from .renderer import simulate_policy, SimulationConfig, ControlPolicy
|
||||
|
||||
__all__ = ["simulate_policy", "SimulationConfig", "ControlPolicy"]
|
||||
|
|
@ -1,71 +0,0 @@
|
|||
from dataclasses import dataclass, fields
|
||||
|
||||
import flax
|
||||
import flax.linen as nn
|
||||
import jax.numpy as jnp
|
||||
import jax.tree_util
|
||||
import numpy as np
|
||||
from flax.linen.initializers import constant, orthogonal
|
||||
|
||||
|
||||
class Network(nn.Module):
|
||||
"""
|
||||
Dummy model only used for testing purposes
|
||||
|
||||
inspired by: https://github.com/vwxyzjn/cleanrl/blob/master/cleanrl/ppo_atari_envpool_xla_jax_scan.py
|
||||
"""
|
||||
|
||||
hidden_dim: int = 195
|
||||
|
||||
@nn.compact
|
||||
def __call__(self, x):
|
||||
x = nn.Dense(self.hidden_dim, kernel_init=orthogonal(np.sqrt(2)), bias_init=constant(0.0))(
|
||||
x
|
||||
)
|
||||
x = nn.relu(x)
|
||||
x = nn.Dense(self.hidden_dim, kernel_init=orthogonal(np.sqrt(2)), bias_init=constant(0.0))(
|
||||
x
|
||||
)
|
||||
x = nn.relu(x)
|
||||
return x
|
||||
|
||||
|
||||
class Critic(nn.Module):
|
||||
@nn.compact
|
||||
def __call__(self, x):
|
||||
return nn.Dense(1, kernel_init=orthogonal(1), bias_init=constant(0.0))(x)
|
||||
|
||||
|
||||
class Actor(nn.Module):
|
||||
action_dim: int
|
||||
|
||||
@nn.compact
|
||||
def __call__(self, x):
|
||||
mean = nn.Dense(self.action_dim, kernel_init=orthogonal(0.01), bias_init=constant(0.0))(x)
|
||||
log_std = self.param("log_std", nn.initializers.zeros, (self.action_dim,))
|
||||
return mean, log_std
|
||||
|
||||
|
||||
@jax.tree_util.register_dataclass
|
||||
@dataclass
|
||||
class AgentParams:
|
||||
network_params: flax.core.FrozenDict
|
||||
actor_params: flax.core.FrozenDict
|
||||
critic_params: flax.core.FrozenDict
|
||||
|
||||
|
||||
@jax.tree_util.register_dataclass
|
||||
@dataclass
|
||||
class Storage:
|
||||
obs: jnp.array
|
||||
actions: jnp.array
|
||||
logprobs: jnp.array
|
||||
dones: jnp.array
|
||||
values: jnp.array
|
||||
advantages: jnp.array
|
||||
returns: jnp.array
|
||||
rewards: jnp.array
|
||||
|
||||
def replace(self, **kwargs) -> "Storage":
|
||||
fs = fields(self)
|
||||
return Storage(**{f.name: kwargs.get(f.name, getattr(self, f.name)) for f in fs})
|
||||
|
|
@ -1,25 +0,0 @@
|
|||
from .DummyAgent import Network, Critic, Actor, AgentParams, Storage
|
||||
from .base import (
|
||||
RLAlgorithm,
|
||||
RLModel,
|
||||
Transition,
|
||||
create_model,
|
||||
register_rl_model,
|
||||
registered_model_types,
|
||||
)
|
||||
from .random_policy_model import RandomPolicyModel
|
||||
|
||||
__all__ = [
|
||||
"RLAlgorithm",
|
||||
"RLModel",
|
||||
"RandomPolicyModel",
|
||||
"Transition",
|
||||
"create_model",
|
||||
"register_rl_model",
|
||||
"registered_model_types",
|
||||
"Network",
|
||||
"Critic",
|
||||
"Actor",
|
||||
"AgentParams",
|
||||
"Storage",
|
||||
]
|
||||
|
|
@ -1,162 +0,0 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Transition:
|
||||
"""A minimal transition container for RL.
|
||||
|
||||
This is intentionally generic because the underlying env state type may be a
|
||||
JAX pytree, a numpy struct, or something library-specific.
|
||||
"""
|
||||
|
||||
obs: Any
|
||||
action: Any
|
||||
reward: float
|
||||
next_obs: Any
|
||||
terminated: bool
|
||||
truncated: bool
|
||||
info: dict[str, Any] | None = None
|
||||
|
||||
|
||||
class RLAlgorithm(ABC):
|
||||
"""Insertable RL algorithm interface."""
|
||||
|
||||
@abstractmethod
|
||||
def select_action(self, *, obs: Any, rng: Any | None = None) -> Any:
|
||||
raise NotImplementedError
|
||||
|
||||
def observe(self, transition: Transition) -> None:
|
||||
"""Optional hook to store transitions."""
|
||||
|
||||
def update(self, *, rng: Any | None = None) -> dict[str, float]:
|
||||
"""Optional hook to run one training update."""
|
||||
|
||||
return {}
|
||||
|
||||
def save(self, path: str) -> None:
|
||||
raise NotImplementedError("Save not implemented")
|
||||
|
||||
def load(self, path: str) -> None:
|
||||
raise NotImplementedError("Load not implemented")
|
||||
|
||||
|
||||
_RL_MODEL_REGISTRY: dict[str, type["RLModel"]] = {}
|
||||
|
||||
|
||||
def registered_model_types() -> list[str]:
|
||||
return sorted(_RL_MODEL_REGISTRY)
|
||||
|
||||
|
||||
def create_model(type_name: str, *, payload: dict[str, Any]) -> "RLModel":
|
||||
model_cls = _RL_MODEL_REGISTRY.get(type_name)
|
||||
if model_cls is None:
|
||||
known = ", ".join(sorted(_RL_MODEL_REGISTRY)) or "<none>"
|
||||
raise ValueError(f"Unknown RLModel type '{type_name}'. Known: {known}")
|
||||
return model_cls.from_payload(payload)
|
||||
|
||||
|
||||
def get_rl_model_registry() -> dict[str, type["RLModel"]]:
|
||||
"""Return a copy of the current RLModel registry.
|
||||
|
||||
The registry is populated by importing concrete model modules that use the
|
||||
`@register_rl_model(...)` decorator.
|
||||
"""
|
||||
|
||||
return dict(_RL_MODEL_REGISTRY)
|
||||
|
||||
|
||||
def register_rl_model(*type_names: str):
|
||||
"""Decorator to register an `RLModel` for generic loading.
|
||||
|
||||
Concrete model modules should apply this decorator, so `base.py` never needs
|
||||
to import concrete models (avoids circular imports).
|
||||
"""
|
||||
|
||||
if not type_names:
|
||||
raise TypeError("register_rl_model() requires at least one type name")
|
||||
|
||||
primary = type_names[0]
|
||||
|
||||
def _decorator(cls: type[RLModel]):
|
||||
for name in type_names:
|
||||
_RL_MODEL_REGISTRY[name] = cls
|
||||
cls.type_name = primary
|
||||
return cls
|
||||
|
||||
return _decorator
|
||||
|
||||
|
||||
class RLModel(ABC):
|
||||
"""Serializable policy/model interface.
|
||||
|
||||
This is the artifact that `train.py` writes and `simulate.py` loads.
|
||||
"""
|
||||
|
||||
# Overwritten by the `@register_rl_model(...)` decorator.
|
||||
type_name: str = "RLModel"
|
||||
|
||||
def reset(self, seed: int | None = None) -> None:
|
||||
"""Optional hook for RNG/stateful models."""
|
||||
|
||||
@abstractmethod
|
||||
def act(self, *, obs: Any | None = None, t: float = 0.0) -> Any:
|
||||
raise NotImplementedError
|
||||
|
||||
def train(self, *, env: Any, num_epochs: int = 1) -> None:
|
||||
"""Optional training hook.
|
||||
|
||||
Many models won't learn; for those this can be a no-op.
|
||||
"""
|
||||
|
||||
_ = (env, num_epochs)
|
||||
|
||||
def to_payload(self) -> dict[str, Any]:
|
||||
"""Return JSON-serializable model parameters."""
|
||||
|
||||
return {}
|
||||
|
||||
@classmethod
|
||||
def from_payload(cls, payload: dict[str, Any]) -> "RLModel":
|
||||
"""Reconstruct a model from `to_payload()` output."""
|
||||
|
||||
return cls(**payload) # type: ignore[arg-type]
|
||||
|
||||
def save(self, path: str | Path) -> Path:
|
||||
out = Path(path)
|
||||
out.parent.mkdir(parents=True, exist_ok=True)
|
||||
doc = {
|
||||
"type": self.type_name,
|
||||
"version": 1,
|
||||
"payload": self.to_payload(),
|
||||
}
|
||||
out.write_text(json.dumps(doc, indent=2, sort_keys=True) + "\n")
|
||||
return out
|
||||
|
||||
@classmethod
|
||||
def load(cls, path: str | Path) -> "RLModel":
|
||||
p = Path(path)
|
||||
doc = json.loads(p.read_text())
|
||||
|
||||
type_name = doc.get("type")
|
||||
if not isinstance(type_name, str):
|
||||
raise ValueError("Model artifact missing string field 'type'")
|
||||
|
||||
model_cls = _RL_MODEL_REGISTRY.get(type_name)
|
||||
if model_cls is None:
|
||||
known = ", ".join(sorted(_RL_MODEL_REGISTRY)) or "<none>"
|
||||
raise ValueError(f"Unknown RLModel type '{type_name}'. Known: {known}")
|
||||
|
||||
payload = doc.get("payload")
|
||||
# Backward compatibility: older artifacts stored fields at top-level.
|
||||
if payload is None:
|
||||
payload = {k: v for k, v in doc.items() if k not in ("type", "version")}
|
||||
if not isinstance(payload, dict):
|
||||
raise ValueError("Model artifact field 'payload' must be an object")
|
||||
|
||||
return model_cls.from_payload(payload)
|
||||
|
|
@ -1,52 +0,0 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import numpy as np
|
||||
|
||||
from .base import RLModel, register_rl_model
|
||||
|
||||
|
||||
@register_rl_model("random")
|
||||
@dataclass(slots=True)
|
||||
class RandomPolicyModel(RLModel):
|
||||
"""A minimal, serializable policy model that outputs random controls.
|
||||
|
||||
This is intentionally *not* a learning algorithm yet. It exists so we can:
|
||||
- produce a stable model artifact from `train.py`
|
||||
- load that artifact in `simulate.py`
|
||||
- drive the MuJoCo viewer with the model's actions
|
||||
"""
|
||||
|
||||
nu: int = 0
|
||||
seed: int = 0
|
||||
ctrl_noise_scale: float = 0.5
|
||||
|
||||
_rng: np.random.RandomState = field(init=False, repr=False)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self.reset(self.seed)
|
||||
|
||||
def reset(self, seed: int | None = None) -> None:
|
||||
if seed is not None:
|
||||
self.seed = int(seed)
|
||||
self._rng = np.random.RandomState(self.seed)
|
||||
|
||||
def act(self, *, obs: np.ndarray | None = None, t: float = 0.0) -> np.ndarray:
|
||||
if self.nu <= 0:
|
||||
return np.zeros((0,), dtype=np.float32)
|
||||
ctrl = self.ctrl_noise_scale * self._rng.randn(self.nu)
|
||||
return ctrl.astype(np.float32)
|
||||
|
||||
def to_payload(self) -> dict[str, object]:
|
||||
return {
|
||||
"seed": int(self.seed),
|
||||
"ctrl_noise_scale": float(self.ctrl_noise_scale),
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_payload(cls, payload: dict[str, object]) -> RandomPolicyModel:
|
||||
return cls(
|
||||
seed=int(payload.get("seed", 0)),
|
||||
ctrl_noise_scale=float(payload.get("ctrl_noise_scale", 0.5)),
|
||||
)
|
||||
536
src/brittle_star_project/trainers/PPOTrainer.py
Normal file
536
src/brittle_star_project/trainers/PPOTrainer.py
Normal file
|
|
@ -0,0 +1,536 @@
|
|||
import datetime
|
||||
import random
|
||||
import sys
|
||||
import time
|
||||
from dataclasses import asdict, dataclass
|
||||
from functools import partial
|
||||
from typing import Any
|
||||
|
||||
import flax
|
||||
import jax
|
||||
import jax.numpy as jnp
|
||||
import numpy as np
|
||||
import optax
|
||||
import tqdm
|
||||
from flax.training.train_state import TrainState
|
||||
from torch.utils.tensorboard import SummaryWriter
|
||||
|
||||
from brittle_star_project.dataclasses import EpisodeStatistics, PPOArgs
|
||||
from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
|
||||
from MLPs.mlps import (
|
||||
Actor,
|
||||
AgentParams,
|
||||
GenericDenseLayersWithActivation,
|
||||
OneDenseLayerMLP,
|
||||
Storage,
|
||||
)
|
||||
from ppo import PPO
|
||||
|
||||
|
||||
@jax.jit
|
||||
def _linear_schedule(count, minibatch_count, update_epochs, num_iterations, learning_rate):
|
||||
frac = 1.0 - (count // (minibatch_count * update_epochs)) / num_iterations
|
||||
return learning_rate * frac
|
||||
|
||||
|
||||
@jax.jit
|
||||
def _convert_obs_dict_to_array(obs_dict: dict) -> jnp.ndarray:
|
||||
return jax.vmap(lambda o: jnp.concatenate([v.flatten() for v in o.values() if v.size > 0]))(
|
||||
obs_dict
|
||||
)
|
||||
|
||||
|
||||
# removed jit: used in _rollout_jit, so will be compiled with _rollout_jit
|
||||
def _get_action_and_value_noise(
|
||||
sensor: GenericDenseLayersWithActivation,
|
||||
feature_extractor: GenericDenseLayersWithActivation,
|
||||
actor: Actor,
|
||||
critic: OneDenseLayerMLP,
|
||||
agent_state: TrainState,
|
||||
next_obs: jnp.ndarray,
|
||||
key: jax.random.PRNGKey,
|
||||
):
|
||||
hidden = sensor.apply(agent_state.params["sensor_params"], next_obs)
|
||||
hidden_critic = feature_extractor.apply(
|
||||
agent_state.params["feature_extractor_params"], next_obs
|
||||
)
|
||||
|
||||
# Continuous actions: sample from a Gaussian parameterized by the actor
|
||||
mean, log_std = actor.apply(agent_state.params["actor_params"], hidden)
|
||||
key, subkey = jax.random.split(key)
|
||||
noise = jax.random.normal(subkey, shape=mean.shape)
|
||||
std = jnp.exp(log_std)
|
||||
action = mean + noise * std
|
||||
logprob = -0.5 * (((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1)
|
||||
value = critic.apply(agent_state.params["critic_params"], hidden_critic)
|
||||
return action, logprob, value.squeeze(-1), key
|
||||
|
||||
|
||||
# removed jit: used in _rollout_jit, so will be compiled with _rollout_jit
|
||||
def _step_once(
|
||||
carry,
|
||||
_,
|
||||
env_step_fn,
|
||||
sensor: GenericDenseLayersWithActivation,
|
||||
feature_extractor: GenericDenseLayersWithActivation,
|
||||
actor: Actor,
|
||||
critic: OneDenseLayerMLP,
|
||||
):
|
||||
agent_state, episode_stats, obs, done, key, env_state = carry
|
||||
action, logprob, value, key = _get_action_and_value_noise(
|
||||
sensor, feature_extractor, actor, critic, agent_state, obs, key
|
||||
)
|
||||
|
||||
episode_stats, env_state, (next_obs, reward, next_done) = env_step_fn(
|
||||
episode_stats, env_state, action
|
||||
)
|
||||
|
||||
storage = Storage(
|
||||
obs=obs,
|
||||
actions=action,
|
||||
logprobs=logprob,
|
||||
dones=done,
|
||||
values=value,
|
||||
rewards=reward,
|
||||
returns=jnp.zeros_like(reward),
|
||||
advantages=jnp.zeros_like(reward),
|
||||
)
|
||||
return (agent_state, episode_stats, next_obs, next_done, key, env_state), storage
|
||||
|
||||
|
||||
# removed jit: used in _rollout_jit, so will be compiled with _rollout_jit
|
||||
def _step_env_wrapped(episode_stats, env_state, action, env_step_fn):
|
||||
next_env_state = env_step_fn(env_state, action)
|
||||
|
||||
# Extract per-environment signals from the state object
|
||||
reward = next_env_state.reward # (num_envs,)
|
||||
terminated = next_env_state.terminated # (num_envs,)
|
||||
truncated = next_env_state.truncated # (num_envs,)
|
||||
done = terminated | truncated # (num_envs,)
|
||||
|
||||
new_episode_return = episode_stats.episode_returns + reward
|
||||
new_episode_length = episode_stats.episode_lengths + 1
|
||||
|
||||
episode_stats = episode_stats.replace(
|
||||
episode_returns=new_episode_return * (1 - done),
|
||||
episode_lengths=new_episode_length * (1 - done),
|
||||
returned_episode_returns=jnp.where(
|
||||
done, new_episode_return, episode_stats.returned_episode_returns
|
||||
),
|
||||
returned_episode_lengths=jnp.where(
|
||||
done, new_episode_length, episode_stats.returned_episode_lengths
|
||||
),
|
||||
)
|
||||
return (
|
||||
episode_stats,
|
||||
next_env_state,
|
||||
(_convert_obs_dict_to_array(next_env_state.observations), reward, done),
|
||||
)
|
||||
|
||||
|
||||
# jit applied in wrapper method self._rollout_jit using partial
|
||||
def _rollout_jit(
|
||||
agent_state,
|
||||
episode_stats,
|
||||
env_state,
|
||||
next_obs,
|
||||
next_done,
|
||||
key,
|
||||
max_steps,
|
||||
step_env_fn,
|
||||
sensor: GenericDenseLayersWithActivation,
|
||||
feature_extractor: GenericDenseLayersWithActivation,
|
||||
actor: Actor,
|
||||
critic: OneDenseLayerMLP,
|
||||
):
|
||||
(agent_state, episode_stats, next_obs, next_done, key, env_state), storage = jax.lax.scan(
|
||||
partial(
|
||||
_step_once,
|
||||
sensor=sensor,
|
||||
feature_extractor=feature_extractor,
|
||||
actor=actor,
|
||||
critic=critic,
|
||||
env_step_fn=step_env_fn,
|
||||
),
|
||||
(agent_state, episode_stats, next_obs, next_done, key, env_state),
|
||||
(),
|
||||
max_steps,
|
||||
)
|
||||
return agent_state, episode_stats, next_obs, next_done, storage, key, env_state
|
||||
|
||||
|
||||
# removed jit: used in _compute_gae_jit, so will be compiled with _compute_gae_jit
|
||||
def _compute_gae_once(carry, inp, gamma, gae_lambda):
|
||||
advantages = carry
|
||||
nextdone, nextvalues, curvalues, reward = inp
|
||||
nextnonterminal = 1.0 - nextdone
|
||||
delta = reward + gamma * nextvalues * nextnonterminal - curvalues
|
||||
advantages = delta + gamma * gae_lambda * nextnonterminal * advantages
|
||||
return advantages, advantages
|
||||
|
||||
|
||||
# jit applied on partial-wrapped wrapper method self._compute_gae_jit
|
||||
def _compute_gae_jit(
|
||||
agent_state, storage, next_obs, next_done, gamma, gae_lambda, num_envs, sensor, critic
|
||||
):
|
||||
next_value = critic.apply(
|
||||
agent_state.params["critic_params"],
|
||||
sensor.apply(agent_state.params["sensor_params"], next_obs),
|
||||
).squeeze(-1)
|
||||
|
||||
advantages = jnp.zeros((num_envs,))
|
||||
dones = jnp.concatenate([storage.dones, next_done[None, :]], axis=0)
|
||||
values = jnp.concatenate([storage.values, next_value[None, :]], axis=0)
|
||||
_, advantages = jax.lax.scan(
|
||||
partial(_compute_gae_once, gamma=gamma, gae_lambda=gae_lambda),
|
||||
advantages,
|
||||
(dones[1:], values[1:], values[:-1], storage.rewards),
|
||||
reverse=True,
|
||||
)
|
||||
return storage.replace(advantages=advantages, returns=advantages + storage.values)
|
||||
|
||||
|
||||
@dataclass
|
||||
class LossInfo:
|
||||
# todo: better typing
|
||||
loss: Any
|
||||
pg_loss: Any
|
||||
v_loss: Any
|
||||
entropy_loss: Any
|
||||
approx_kl: Any
|
||||
avg_episodic_return: Any
|
||||
|
||||
|
||||
class PPOTrainer:
|
||||
def __init__(self, args: PPOArgs, env: BrittleStarJaxEnvWrapper, run_dir: str, run_name: str):
|
||||
self.args = args
|
||||
self.env = env
|
||||
self.run_dir = run_dir
|
||||
self.run_name = run_name
|
||||
self.writer = SummaryWriter(self.run_dir)
|
||||
|
||||
self.key = jax.random.PRNGKey(args.seed)
|
||||
|
||||
self.sensor, self.feature_extractor, self.actor, self.critic = self._init_agent()
|
||||
self.sensor.apply = jax.jit(self.sensor.apply)
|
||||
self.feature_extractor.apply = jax.jit(self.feature_extractor.apply)
|
||||
self.actor.apply = jax.jit(self.actor.apply)
|
||||
self.critic.apply = jax.jit(self.critic.apply)
|
||||
|
||||
self._rollout_jit = jax.jit(
|
||||
partial(
|
||||
_rollout_jit,
|
||||
max_steps=self.args.num_steps,
|
||||
step_env_fn=partial(_step_env_wrapped, env_step_fn=self.env.step),
|
||||
sensor=self.sensor,
|
||||
feature_extractor=self.feature_extractor,
|
||||
actor=self.actor,
|
||||
critic=self.critic,
|
||||
)
|
||||
)
|
||||
self._compute_gae_jit = jax.jit(
|
||||
partial(
|
||||
_compute_gae_jit,
|
||||
num_envs=self.args.num_envs,
|
||||
gamma=self.args.gamma,
|
||||
gae_lambda=self.args.gae_lambda,
|
||||
sensor=self.sensor,
|
||||
critic=self.critic,
|
||||
)
|
||||
)
|
||||
|
||||
self._ppo = PPO(self.args, self.sensor, self.actor, self.critic, self.feature_extractor)
|
||||
|
||||
self.agent_state = self._init_agent_state()
|
||||
|
||||
self.episode_stats = self._init_episode_stats()
|
||||
|
||||
self._init_random()
|
||||
|
||||
def _init_random(self, log: bool = True):
|
||||
if log:
|
||||
print(f"[RANDOM]: Setting random seed to {self.args.seed}")
|
||||
|
||||
random.seed(self.args.seed)
|
||||
np.random.seed(self.args.seed)
|
||||
|
||||
def _init_agent(self, log: bool = True):
|
||||
if log:
|
||||
print("[AGENT]: Initializing agent...")
|
||||
|
||||
sensor = GenericDenseLayersWithActivation()
|
||||
feature_extractor = GenericDenseLayersWithActivation()
|
||||
actor = Actor(
|
||||
action_dim=self.env.single_action_space.shape[0]
|
||||
) # continuous actions for MJX
|
||||
critic = OneDenseLayerMLP()
|
||||
# messenger = OneDenseLayerMLP()
|
||||
return sensor, feature_extractor, actor, critic
|
||||
|
||||
def _init_agent_state(self, log: bool = True) -> TrainState:
|
||||
if log:
|
||||
print("[AGENT STATE]: Initializing agent state...")
|
||||
|
||||
self.key, sensor_key, actor_key, critic_key, feature_extractor_key = jax.random.split(
|
||||
self.key, 5
|
||||
)
|
||||
|
||||
sample_obs = jnp.concatenate(
|
||||
[
|
||||
v.flatten()
|
||||
for v in self.env.single_observation_space.sample(
|
||||
rng=jax.random.PRNGKey(0)
|
||||
).values()
|
||||
if v.size > 0
|
||||
]
|
||||
)
|
||||
sensor_params = self.sensor.init(sensor_key, sample_obs)
|
||||
feature_extractor_params = self.feature_extractor.init(feature_extractor_key, sample_obs)
|
||||
actor_params = self.actor.init(actor_key, self.sensor.apply(sensor_params, sample_obs))
|
||||
critic_params = self.critic.init(
|
||||
critic_key, self.feature_extractor.apply(feature_extractor_params, sample_obs)
|
||||
)
|
||||
|
||||
return TrainState.create(
|
||||
apply_fn=None,
|
||||
params=asdict(
|
||||
AgentParams(sensor_params, actor_params, critic_params, feature_extractor_params)
|
||||
),
|
||||
tx=optax.chain(
|
||||
optax.clip_by_global_norm(self.args.max_grad_norm),
|
||||
optax.inject_hyperparams(optax.adam)(
|
||||
learning_rate=partial(
|
||||
_linear_schedule,
|
||||
minibatch_count=self.args.num_minibatches,
|
||||
update_epochs=self.args.update_epochs,
|
||||
num_iterations=self.args.num_iterations,
|
||||
learning_rate=self.args.learning_rate,
|
||||
)
|
||||
if self.args.anneal_lr
|
||||
else self.args.learning_rate,
|
||||
eps=1e-5,
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
def _init_episode_stats(self, log: bool = True) -> EpisodeStatistics:
|
||||
if log:
|
||||
print("[EPISODE STATS]: Initializing episode stats...")
|
||||
|
||||
return EpisodeStatistics(
|
||||
episode_returns=jnp.zeros(self.args.num_envs, dtype=jnp.float32),
|
||||
episode_lengths=jnp.zeros(self.args.num_envs, dtype=jnp.int32),
|
||||
returned_episode_returns=jnp.zeros(self.args.num_envs, jnp.float32),
|
||||
returned_episode_lengths=jnp.zeros(self.args.num_envs, dtype=jnp.int32),
|
||||
)
|
||||
|
||||
def _rollout(self, env_state, next_obs, next_done) -> tuple[Storage, ...]:
|
||||
return self._rollout_jit(
|
||||
self.agent_state,
|
||||
self.episode_stats,
|
||||
env_state,
|
||||
next_obs,
|
||||
next_done,
|
||||
self.key,
|
||||
)
|
||||
|
||||
def _compute_gae(self, storage, next_obs, next_done) -> Storage:
|
||||
return self._compute_gae_jit(
|
||||
self.agent_state,
|
||||
storage,
|
||||
next_obs,
|
||||
next_done,
|
||||
)
|
||||
|
||||
def _log(
|
||||
self,
|
||||
global_step,
|
||||
episode_stats,
|
||||
start_time,
|
||||
iteration_time_start,
|
||||
loss_info,
|
||||
):
|
||||
|
||||
self.writer.add_scalar(
|
||||
"charts/avg_episodic_return", loss_info.avg_episodic_return, global_step
|
||||
)
|
||||
self.writer.add_scalar(
|
||||
"charts/avg_episodic_length",
|
||||
np.mean(jax.device_get(episode_stats.returned_episode_lengths)),
|
||||
global_step,
|
||||
)
|
||||
self.writer.add_scalar(
|
||||
"charts/learning_rate",
|
||||
self.agent_state.opt_state[1].hyperparams["learning_rate"].item(),
|
||||
global_step,
|
||||
)
|
||||
self.writer.add_scalar("losses/value_loss", loss_info.v_loss[-1, -1].item(), global_step)
|
||||
self.writer.add_scalar("losses/policy_loss", loss_info.pg_loss[-1, -1].item(), global_step)
|
||||
self.writer.add_scalar("losses/entropy", loss_info.entropy_loss[-1, -1].item(), global_step)
|
||||
self.writer.add_scalar("losses/approx_kl", loss_info.approx_kl[-1, -1].item(), global_step)
|
||||
self.writer.add_scalar("losses/loss", loss_info.loss[-1, -1].item(), global_step)
|
||||
self.writer.add_scalar(
|
||||
"charts/SPS", int(global_step / (time.time() - start_time)), global_step
|
||||
)
|
||||
self.writer.add_scalar(
|
||||
"charts/SPS_update",
|
||||
int(self.args.num_envs * self.args.num_steps / (time.time() - iteration_time_start)),
|
||||
global_step,
|
||||
)
|
||||
|
||||
def _step(
|
||||
self, env_state, next_obs, next_done, is_tty: bool, iteration: int, log: bool = True
|
||||
) -> tuple:
|
||||
if log and not is_tty and iteration == 1:
|
||||
print(f">>> [HPC] Starting first rollout (JIT): {time.ctime()}", flush=True)
|
||||
|
||||
(
|
||||
self.agent_state,
|
||||
self.episode_stats,
|
||||
next_obs,
|
||||
next_done,
|
||||
storage,
|
||||
self.key,
|
||||
next_env_state,
|
||||
) = self._rollout(env_state, next_obs, next_done)
|
||||
|
||||
if log and not is_tty and iteration == 1:
|
||||
print(f">>> [HPC] First rollout completed: {time.ctime()}", flush=True)
|
||||
|
||||
storage = self._compute_gae(storage, next_obs, next_done)
|
||||
|
||||
if log and not is_tty and iteration == 1:
|
||||
print(f">>> [HPC] Starting first PPO update (JIT): {time.ctime()}", flush=True)
|
||||
|
||||
self.agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, self.key = (
|
||||
self._ppo.update_ppo(self.agent_state, storage, self.key)
|
||||
)
|
||||
|
||||
if log and not is_tty and iteration == 1:
|
||||
print(f">>> [HPC] First PPO update completed: {time.ctime()}", flush=True)
|
||||
|
||||
avg_episodic_return = float(
|
||||
jnp.mean(jax.device_get(self.episode_stats.returned_episode_returns)).item()
|
||||
)
|
||||
|
||||
return (
|
||||
next_env_state,
|
||||
next_obs,
|
||||
next_done,
|
||||
LossInfo(
|
||||
loss=loss,
|
||||
pg_loss=pg_loss,
|
||||
v_loss=v_loss,
|
||||
entropy_loss=entropy_loss,
|
||||
approx_kl=approx_kl,
|
||||
avg_episodic_return=avg_episodic_return,
|
||||
),
|
||||
)
|
||||
|
||||
def _close(self):
|
||||
self.env.close()
|
||||
self.writer.close()
|
||||
|
||||
def _save_model(self, model_path: str, log: bool = True):
|
||||
if log:
|
||||
print(f"[SAVE]: Saving the model to: {model_path}...")
|
||||
|
||||
with open(model_path, "wb") as f:
|
||||
f.write(
|
||||
flax.serialization.to_bytes(
|
||||
[
|
||||
vars(self.args),
|
||||
[
|
||||
self.agent_state.params["sensor_params"],
|
||||
self.agent_state.params["actor_params"],
|
||||
self.agent_state.params["critic_params"],
|
||||
self.agent_state.params["feature_extractor_params"],
|
||||
],
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
def train(self, log: bool = True):
|
||||
"""
|
||||
Train the PPO agent for a specified number of iterations
|
||||
(passed through PPOArgs in constructor).
|
||||
Closes the environment at the end of training.
|
||||
"""
|
||||
if log:
|
||||
print(f"running name: {self.run_name}")
|
||||
|
||||
is_tty = sys.stdout.isatty()
|
||||
if log:
|
||||
print("[TRAIN]: Resetting environment...")
|
||||
|
||||
if not is_tty:
|
||||
print(f">>> [HPC] Initial reset started: {time.ctime()}", flush=True)
|
||||
|
||||
env_state = self.env.reset(seed=self.args.seed)
|
||||
next_obs = _convert_obs_dict_to_array(env_state.observations)
|
||||
next_done = jnp.zeros(self.args.num_envs, dtype=jnp.bool_)
|
||||
|
||||
if log and not is_tty:
|
||||
print(f">>> [HPC] Initial reset completed: {time.ctime()}", flush=True)
|
||||
|
||||
global_step = 0
|
||||
start_time = time.time()
|
||||
|
||||
if self.args.track:
|
||||
import wandb
|
||||
|
||||
if log:
|
||||
print("[TRAIN]: Initializing Weights and Biases...")
|
||||
|
||||
wandb.init(
|
||||
project=self.args.wandb_project_name,
|
||||
entity=self.args.wandb_entity,
|
||||
sync_tensorboard=True,
|
||||
config=vars(self.args),
|
||||
name=self.run_name,
|
||||
save_code=True,
|
||||
)
|
||||
|
||||
if log:
|
||||
print("[TRAIN]: Adding hyperparameters to TensorBoard...")
|
||||
|
||||
self.writer.add_text(
|
||||
"hyperparameters",
|
||||
"|param|value|\n|---|---|\n"
|
||||
+ "\n".join(f"|{k}|{v}|" for k, v in vars(self.args).items()),
|
||||
)
|
||||
|
||||
iter_bar = tqdm.tqdm(
|
||||
range(1, self.args.num_iterations + 1),
|
||||
disable=not is_tty,
|
||||
)
|
||||
for iteration in iter_bar:
|
||||
iteration_time_start = time.time()
|
||||
|
||||
env_state, next_obs, next_done, loss_info = self._step(
|
||||
env_state, next_obs, next_done, is_tty=is_tty, iteration=iteration
|
||||
)
|
||||
|
||||
global_step += self.args.num_steps * self.args.num_envs
|
||||
self._log(global_step, self.episode_stats, start_time, iteration_time_start, loss_info)
|
||||
|
||||
if log and not is_tty:
|
||||
sps = int(global_step / (time.time() - start_time))
|
||||
remaining_steps = self.args.total_timesteps - global_step
|
||||
eta_seconds = int(remaining_steps / sps) if sps > 0 else 0
|
||||
eta_str = str(datetime.timedelta(seconds=eta_seconds))
|
||||
|
||||
print(
|
||||
f"Iteration {iteration}/{self.args.num_iterations} | "
|
||||
f"Step {global_step}/{self.args.total_timesteps} | "
|
||||
f"SPS {sps} | "
|
||||
f"Return {loss_info.avg_episodic_return:.4f} | "
|
||||
f"ETA {eta_str}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
if self.args.save_model:
|
||||
model_path = f"{self.run_dir}/{self.args.exp_name}.cleanrl_model"
|
||||
self._save_model(model_path=model_path)
|
||||
|
||||
self._close()
|
||||
0
src/brittle_star_project/trainers/__init__.py
Normal file
0
src/brittle_star_project/trainers/__init__.py
Normal file
Reference in a new issue