Move experiments sub-package back to corax.jax

This is primarily for helping to migrate other projects
ethanluoyc · Oct 10, 2023 · 96809cb · 96809cb
1 parent fded3c4
commit 96809cb
Show file tree

Hide file tree

Showing 12 changed files with 32 additions and 46 deletions.
diff --git a/corax/jax/experiments/__init__.py b/corax/jax/experiments/__init__.py
@@ -0,0 +1,18 @@
+"""JAX experiment utils."""
+
+from corax.jax.experiments.config import CheckpointingConfig
+from corax.jax.experiments.config import DeprecatedPolicyFactory
+from corax.jax.experiments.config import EvaluatorFactory
+from corax.jax.experiments.config import ExperimentConfig
+from corax.jax.experiments.config import MakeActorFn
+from corax.jax.experiments.config import NetworkFactory
+from corax.jax.experiments.config import OfflineExperimentConfig
+from corax.jax.experiments.config import PolicyFactory
+from corax.jax.experiments.config import default_evaluator_factory
+from corax.jax.experiments.config import make_policy
+from corax.jax.experiments.imitation_experiment import ImitationExperimentConfig
+from corax.jax.experiments.imitation_experiment import run_imitation_experiment
+from corax.jax.experiments.imitation_loop import EpisodeRewarder
+from corax.jax.experiments.imitation_loop import ImitationEnvironmentLoop
+from corax.jax.experiments.run_experiment import run_experiment
+from corax.jax.experiments.run_offline_experiment import run_offline_experiment
diff --git a/...baselines/baselines/experiments/config.py → corax/jax/experiments/config.py b/...baselines/baselines/experiments/config.py → corax/jax/experiments/config.py
diff --git a/...lines/experiments/imitation_experiment.py → ...x/jax/experiments/imitation_experiment.py b/...lines/experiments/imitation_experiment.py → ...x/jax/experiments/imitation_experiment.py
@@ -8,14 +8,14 @@
 import reverb
 
 import corax
-from baselines import experiments
-from baselines.experiments import config
-from baselines.experiments import imitation_loop
-from baselines.experiments import savers
-from baselines.experiments.run_experiment import _disable_insert_blocking
-from baselines.experiments.run_experiment import _LearningActor
 from corax import specs as env_specs
+from corax.jax import experiments
 from corax.jax import utils
+from corax.jax.experiments import config
+from corax.jax.experiments import imitation_loop
+from corax.jax.experiments import savers
+from corax.jax.experiments.run_experiment import _disable_insert_blocking
+from corax.jax.experiments.run_experiment import _LearningActor
 from corax.utils import counting
 
 Config = experiments.ExperimentConfig

diff --git a/...s/baselines/experiments/imitation_loop.py → corax/jax/experiments/imitation_loop.py b/...s/baselines/experiments/imitation_loop.py → corax/jax/experiments/imitation_loop.py
diff --git a/...elines/experiments/imitation_loop_test.py → corax/jax/experiments/imitation_loop_test.py b/...elines/experiments/imitation_loop_test.py → corax/jax/experiments/imitation_loop_test.py
@@ -1,8 +1,8 @@
 import numpy as np
 from absl.testing import absltest
 
-from baselines.experiments import imitation_loop
 from corax import specs
+from corax.jax.experiments import imitation_loop
 from corax.testing import fakes
 
 

diff --git a/...s/baselines/experiments/run_experiment.py → corax/jax/experiments/run_experiment.py b/...s/baselines/experiments/run_experiment.py → corax/jax/experiments/run_experiment.py
@@ -24,12 +24,12 @@
 import reverb
 
 import corax
-from baselines.experiments import config
-from baselines.experiments import savers
 from corax import core
 from corax import specs
 from corax import types
 from corax.jax import utils
+from corax.jax.experiments import config
+from corax.jax.experiments import savers
 from corax.utils import counting
 
 

diff --git a/...nes/experiments/run_offline_experiment.py → ...jax/experiments/run_offline_experiment.py b/...nes/experiments/run_offline_experiment.py → ...jax/experiments/run_offline_experiment.py
@@ -18,9 +18,9 @@
 import jax
 
 import corax
-from baselines.experiments import config
-from baselines.experiments import savers
 from corax import specs
+from corax.jax.experiments import config
+from corax.jax.experiments import savers
 from corax.utils import counting
 
 

diff --git a/...baselines/baselines/experiments/savers.py → corax/jax/experiments/savers.py b/...baselines/baselines/experiments/savers.py → corax/jax/experiments/savers.py
diff --git a/projects/baselines/baselines/experiments/__init__.py b/projects/baselines/baselines/experiments/__init__.py
diff --git a/projects/baselines/baselines/rlpd/main.py b/projects/baselines/baselines/rlpd/main.py
@@ -5,9 +5,9 @@
 
 from baselines import d4rl_utils
 from baselines import experiment_utils
-from baselines import experiments
 from corax.agents.jax import redq
 from corax.datasets import tfds
+from corax.jax import experiments
 
 _CONFIG = config_flags.DEFINE_config_file(
     "config", None, "File path to the training configuration."

diff --git a/projects/baselines/baselines/tdmpc/imitation.py b/projects/baselines/baselines/tdmpc/imitation.py
@@ -20,10 +20,10 @@
 from ml_collections import config_flags
 
 from baselines import experiment_utils
-from baselines import experiments
 from baselines.tdmpc import rewarder as rewarder_lib
 from corax import types
 from corax.agents.jax import tdmpc
+from corax.jax import experiments
 from corax.jax import variable_utils
 
 os.environ["MUJOCO_GL"] = "egl"

diff --git a/projects/baselines/baselines/tdmpc/main.py b/projects/baselines/baselines/tdmpc/main.py
@@ -16,8 +16,8 @@
 from ml_collections import config_flags
 
 from baselines import experiment_utils
-from baselines import experiments
 from corax.agents.jax import tdmpc
+from corax.jax import experiments
 
 os.environ["MUJOCO_GL"] = "egl"  # noqa