fit new emei api

.gitignore

@@ -138,6 +138,4 @@ cython_debug/
 .vscode
 .idea
 
-/cmrl.egg-info/
 /exp/
-/stable-baselines3/

CONTRIBUTING.md

@@ -43,7 +43,7 @@ Emoji | Description
 :art: `:art:` | When you improved / added assets like themes.
 :rocket: `:rocket:` | When you improved performance.
 :memo: `:memo:` | When you wrote documentation.
-:beetle: `:beetle:` | When you fixed a bug.
+:bug: `:bug:` | When you fixed a bug.
 :twisted_rightwards_arrows: `:twisted_rightwards_arrows:` | When you merged a branch.
 :fire: `:fire:` | When you removed something.
 :truck: `:truck:` | When you moved / renamed something.

README.md

@@ -1,4 +1,4 @@
-![](/img/cmrl_logo.png)
+![](/docs/cmrl_logo.png)
 
 # Causal-MBRL
 
@@ -10,7 +10,7 @@
 <a href="https://www.python.org/downloads/release/python-380/"><img src="https://img.shields.io/badge/python-3.8-brightgreen"></a>
 
 `cmrl`(short for `Causal-MBRL`) is a toolbox for facilitating the development of Causal Model-based Reinforcement
-learning algorithms. It use [Stable-Baselines3](https://github.com/DLR-RM/stable-baselines3) as model-free engine and
+learning algorithms. It uses [Stable-Baselines3](https://github.com/DLR-RM/stable-baselines3) as model-free engine and
 allows flexible use of causal models.
 
 `cmrl` is inspired by [MBRL-Lib](https://github.com/facebookresearch/mbrl-lib). Unlike MBRL-Lib, `cmrl` focuses on the
@@ -111,18 +111,44 @@ cd causal-mbrl
 # create conda env
 conda create -n cmrl python=3.8
 conda activate cmrl
+# install torch
+conda install pytorch -c pytorch
 # install cmrl and its dependent packages
 pip install -e .
 ```
 
-If there is no `cuda` in your device, it's convenient to install `cuda` and `pytorch` from conda directly (refer
-to [pytorch](https://pytorch.org/get-started/locally/)):
+for pytorch
 
-````shell
-# for example, in the case of cuda=11.3
-conda install pytorch cudatoolkit=11.3 -c pytorch
-````
+```shell
+# for MacOS
+conda install pytorch -c pytorch
+# for Linux
+conda install pytorch pytorch-cuda=11.6 -c pytorch -c nvidia
+```
+
+for KCIT and RCIT
+
+```shell
+conda install -c conda-forge r-base
+conda install -c conda-forge r-devtools
+R
+```
 
+```shell
+# Install the RCIT from Github. 
+install.packages("devtools")
+library(devtools)
+install_github("ericstrobl/RCIT")
+library(RCIT)
+
+# Install R libraries for RCIT
+install.packages("MASS")
+install.packages("momentchi2")
+install.packages("devtools")
+
+# test RCIT
+RCIT(rnorm(1000),rnorm(1000),rnorm(1000))
+```
 ## install using pip
 
 coming soon.

cmrl/algorithms/__init__.py

@@ -1,4 +1,4 @@
-from cmrl.algorithms.offline import mopo
-from cmrl.algorithms.offline import off_dyna
-from cmrl.algorithms.online import mbpo
-from cmrl.algorithms.online import on_dyna
+from cmrl.algorithms.off_dyna import OfflineDyna
+from cmrl.algorithms.mopo import MOPO
+from cmrl.algorithms.on_dyna import OnlineDyna
+from cmrl.algorithms.mbpo import MBPO

cmrl/algorithms/base_algorithm.py

@@ -0,0 +1,116 @@
+import os
+from typing import Optional
+from functools import partial
+
+import numpy as np
+import torch
+from omegaconf import DictConfig, OmegaConf
+from stable_baselines3.common.buffers import ReplayBuffer
+from stable_baselines3.common.callbacks import BaseCallback
+import wandb
+
+from cmrl.models.fake_env import VecFakeEnv
+from cmrl.sb3_extension.logger import configure as logger_configure
+from cmrl.sb3_extension.eval_callback import EvalCallback
+from cmrl.utils.creator import create_dynamics, create_agent
+from cmrl.utils.env import make_env
+
+
+class BaseAlgorithm:
+    def __init__(
+        self,
+        cfg: DictConfig,
+        work_dir: Optional[str] = None,
+    ):
+        self.cfg = cfg
+        self.work_dir = work_dir or os.getcwd()
+
+        self.env, fns = make_env(self.cfg)
+        self.reward_fn, self.termination_fn, self.get_init_obs_fn, self.obs2state_fn, self.state2obs_fn = fns
+
+        self.eval_env, *_ = make_env(self.cfg)
+        np.random.seed(self.cfg.seed)
+        torch.manual_seed(self.cfg.seed)
+
+        format_strings = ["tensorboard", "multi_csv"]
+        if self.cfg.verbose:
+            format_strings += ["stdout"]
+        self.logger = logger_configure("log", format_strings)
+
+        if cfg.wandb:
+            wandb.init(
+                project="causal-mbrl",
+                group=cfg.exp_name,
+                config=OmegaConf.to_container(cfg, resolve=True),
+                sync_tensorboard=True,
+            )
+
+        # create ``cmrl`` dynamics
+        self.dynamics = create_dynamics(
+            self.cfg, self.env.state_space, self.env.action_space, self.obs2state_fn, self.state2obs_fn, logger=self.logger
+        )
+
+        if self.cfg.transition.name == "oracle_transition":
+            graph = self.env.get_transition_graph() if self.cfg.transition.oracle == "truth" else None
+            self.dynamics.transition.set_oracle_graph(graph)
+        if self.cfg.reward_mech.learn and not self.cfg.reward_mech.name == "oracle_reward_mech":
+            graph = self.env.get_reward_mech_graph() if self.cfg.transition.oracle == "truth" else None
+            self.dynamics.reward_mech.set_oracle_graph(graph)
+        if self.cfg.termination_mech.learn and not self.cfg.termination_mech.name == "oracle_termination_mech":
+            graph = self.env.get_termination_mech_graph() if self.cfg.transition.oracle == "truth" else None
+            self.dynamics.termination_mech.set_oracle_graph(graph)
+
+        # create sb3's replay buffer for real offline data
+        self.real_replay_buffer = ReplayBuffer(
+            cfg.task.num_steps,
+            self.env.observation_space,
+            self.env.action_space,
+            self.cfg.device,
+            handle_timeout_termination=False,
+        )
+
+        self.partial_fake_env = partial(
+            VecFakeEnv,
+            self.cfg.algorithm.num_envs,
+            self.env.state_space,
+            self.env.action_space,
+            self.dynamics,
+            self.reward_fn,
+            self.termination_fn,
+            self.get_init_obs_fn,
+            self.real_replay_buffer,
+            penalty_coeff=self.cfg.task.penalty_coeff,
+            logger=self.logger,
+        )
+        self.agent = create_agent(self.cfg, self.fake_env, self.logger)
+
+    @property
+    def fake_env(self) -> VecFakeEnv:
+        return self.partial_fake_env(
+            deterministic=self.cfg.algorithm.deterministic,
+            max_episode_steps=self.env.spec.max_episode_steps,
+            branch_rollout=False,
+        )
+
+    @property
+    def callback(self) -> BaseCallback:
+        fake_eval_env = self.partial_fake_env(
+            deterministic=True, max_episode_steps=self.env.spec.max_episode_steps, branch_rollout=False
+        )
+        return EvalCallback(
+            self.eval_env,
+            fake_eval_env,
+            n_eval_episodes=self.cfg.task.n_eval_episodes,
+            best_model_save_path="./",
+            eval_freq=self.cfg.task.eval_freq,
+            deterministic=True,
+            render=False,
+        )
+
+    def learn(self):
+        self._setup_learn()
+
+        self.agent.learn(total_timesteps=self.cfg.task.num_steps, callback=self.callback)
+
+    def _setup_learn(self):
+        pass