Add temporary testing files

Author: yardenas

config/hardware/2080_rtx.yaml

@@ -0,0 +1,6 @@
+# @package _global_
+
+hydra:
+  launcher:
+    additional_parameters: { "gpus": "rtx_2080_ti:1", "account": "ls_krausea" }
+

config/hardware/3090_rtx.yaml

@@ -0,0 +1,6 @@
+# @package _global_
+
+hydra:
+  launcher:
+    additional_parameters: { "gpus": "rtx_3090:1", "account": "ls_krausea" }
+

config/hardware/4090_rtx.yaml

@@ -0,0 +1,6 @@
+# @package _global_
+
+hydra:
+  launcher:
+    additional_parameters: { "gpus": "rtx_4090:1", "account": "ls_krausea" }
+

config/hydra/launcher/slurm.yaml

@@ -0,0 +1,16 @@
+submitit_folder: ${hydra.sweep.dir}/.submitit/%j
+timeout_min: 30
+cpus_per_task: 10
+tasks_per_node: 1
+mem_gb: null
+nodes: 1
+name: ${hydra.job.name}
+_target_: hydra_plugins.hydra_submitit_launcher.submitit_launcher.SlurmLauncher
+mem_per_gpu: null
+mem_per_cpu: 10240
+account: ls_krausea
+additional_parameters: {"gpus": "rtx_4090:1", "account": "ls_krausea"}
+array_parallelism: 256
+max_num_timeout: 100
+setup:
+  - '#SBATCH --requeue'

config/train_brax.yaml

@@ -0,0 +1,29 @@
+defaults:
+  - _self_
+
+hydra:
+  run:
+    dir: ${log_dir}/${now:%Y-%m-%d}/${now:%H-%M-%S}
+  sweep:
+    dir: ${log_dir}/${hydra.job.name}
+    subdir: ${hydra.job.override_dirname}/seed=${training.seed}
+  job:
+    config:
+      override_dirname:
+        exclude_keys:
+          - log_dir
+          - training.seed
+          - wandb
+    chdir: true
+
+
+wandb:
+  group: null
+  notes: null
+  name: ${hydra:job.override_dirname}
+
+jit: true
+
+training:
+  seed: 0
+  render: true

train_brax.py

@@ -0,0 +1,208 @@
+import functools
+import logging
+import os
+
+import hydra
+import jax
+from mujoco_playground import registry, wrapper
+from omegaconf import OmegaConf
+import jax.numpy as jp
+
+
+import logging
+import os
+from typing import Any
+
+import numpy as np
+import omegaconf
+from numpy import typing as npt
+from omegaconf import DictConfig
+from omegaconf.errors import InterpolationKeyError
+
+_LOG = logging.getLogger(__name__)
+
+
+class WeightAndBiasesWriter:
+    def __init__(self, config: DictConfig):
+        import wandb
+
+        try:
+            name = config.wandb.name
+        except InterpolationKeyError:
+            name = None
+        config.wandb.name = name
+        config_dict = omegaconf.OmegaConf.to_container(config, resolve=True)
+        assert isinstance(config_dict, dict)
+        wandb.init(project="ss2r", resume=True, config=config_dict, **config.wandb)
+        self._handle = wandb
+
+    def log(self, summary: dict[str, float], step: int):
+        self._handle.log(summary, step=step)
+
+    def log_video(
+        self,
+        images: npt.ArrayLike,
+        step: int,
+        name: str = "policy",
+        fps: int | float = 30,
+    ):
+        self._handle.log(
+            {
+                name: self._handle.Video(
+                    np.array(images, copy=False),
+                    fps=int(fps),
+                    caption=name,
+                )
+            },
+            step=step,
+        )
+
+    def log_artifact(
+        self,
+        path: str,
+        type: str,
+        name: str | None = None,
+        description: str | None = None,
+        metadata: dict[str, Any] | None = None,
+    ):
+        if name is None:
+            name = self._handle.run.id
+        if metadata is None:
+            metadata = dict(self._handle.config)
+        artifact = self._handle.Artifact(name, type, description, metadata)
+        artifact.add_file(path)
+        self._handle.log_artifact(artifact, aliases=[self._handle.run.id])
+
+
+def get_state_path() -> str:
+    log_path = os.getcwd()
+    return log_path
+
+
+env_name = "QuadrupedRun"
+env = registry.load(env_name)
+env_cfg = registry.get_default_config(env_name)
+eval_env = registry.load(env_name, config=env_cfg)
+agent_name = "PPO"
+
+
+def get_ppo_train_fn():
+    from brax.training.agents.ppo import networks as ppo_networks
+    from brax.training.agents.ppo import train as ppo
+    from mujoco_playground.config import locomotion_params
+
+    ppo_params = locomotion_params.brax_ppo_config(env_name)
+    ppo_training_params = dict(ppo_params)
+    network_factory = ppo_networks.make_ppo_networks
+    if "network_factory" in ppo_params:
+        del ppo_training_params["network_factory"]
+        network_factory = functools.partial(
+            ppo_networks.make_ppo_networks, **ppo_params.network_factory
+        )
+    train_fn = functools.partial(
+        ppo.train,
+        **dict(ppo_training_params),
+        network_factory=network_factory,
+    )
+    return train_fn
+
+
+class Counter:
+    def __init__(self):
+        self.count = 0
+
+
+def report(logger, step, num_steps, metrics):
+    metrics = {k: float(v) for k, v in metrics.items()}
+    logger.log(metrics, num_steps)
+    step.count = num_steps
+
+
+@functools.partial(jax.jit, static_argnames=("env", "policy", "steps"))
+def rollout(
+    env,
+    policy,
+    steps,
+    rng,
+    state,
+):
+    def f(carry, _):
+        state, current_key = carry
+        current_key, next_key = jax.random.split(current_key)
+        action, _ = policy(state.obs, current_key)
+        nstate = env.step(
+            state,
+            action,
+        )
+        return (nstate, next_key), nstate
+
+    (final_state, _), data = jax.lax.scan(f, (state, rng), (), length=steps)
+    return final_state, data
+
+
+def pytrees_unstack(pytree):
+    leaves, treedef = jax.tree_flatten(pytree)
+    n_trees = leaves[0].shape[0]
+    new_leaves = [[] for _ in range(n_trees)]
+    for leaf in leaves:
+        for i in range(n_trees):
+            new_leaves[i].append(leaf[i])
+    new_trees = [treedef.unflatten(leaf) for leaf in new_leaves]
+    return new_trees
+
+
+def render(env, policy, steps, rng, camera=None):
+    state = env.reset(rng)
+    state = jax.tree_map(lambda x: x[:5], state)
+    orig_model = env._mjx_model
+    if hasattr(env, "_randomized_models"):
+        render_env = env
+        model = jax.tree_map(
+            lambda x, ax: jp.take(x, jp.arange(5), axis=ax) if ax is not None else x,
+            env._randomized_models,
+            env._in_axes,
+        )
+        render_env._randomized_models = model
+    else:
+        render_env = env
+    _, trajectory = rollout(render_env, policy, steps, rng[0], state)
+    env._mjx_model = orig_model
+    videos = []
+    for i in range(5):
+        ep_trajectory = jax.tree_map(lambda x: x[:, i], trajectory)
+        ep_trajectory = pytrees_unstack(ep_trajectory)
+        video = env.render(ep_trajectory, camera=camera)
+        videos.append(video)
+    return np.asarray(videos).transpose(0, 1, 4, 2, 3)
+
+
+@hydra.main(version_base=None, config_path="ss2r/configs", config_name="train_brax")
+def main(cfg):
+    _LOG.info(
+        f"Setting up experiment with the following configuration: "
+        f"\n{OmegaConf.to_yaml(cfg)}"
+    )
+    logger = WeightAndBiasesWriter(cfg)
+    train_fn = get_ppo_train_fn()
+    steps = Counter()
+    with jax.disable_jit(not cfg.jit):
+        make_policy, params, _ = train_fn(
+            environment=env,
+            eval_env=eval_env,
+            wrap_env_fn=wrapper.wrap_for_brax_training,
+            progress_fn=functools.partial(report, logger, steps),
+        )
+    if cfg.training.render:
+        rng = jax.random.split(jax.random.PRNGKey(cfg.training.seed), 128)
+        video = render(
+            eval_env,
+            make_policy(params, deterministic=True),
+            1000,
+            rng,
+        )
+        logger.log_video(video, steps.count, "eval/video")
+    _LOG.info("Done training.")
+
+
+if __name__ == "__main__":
+    main()