From 1d67431d4a1fb5032d78a0b3d235c494286f520b Mon Sep 17 00:00:00 2001
From: Benjamin Bolte <ben@bolte.cc>
Date: Wed, 16 Oct 2024 13:06:08 -0700
Subject: [PATCH] more improvements

---
 README.md        | 12 ++++++++++++
 minppo/cli.py    |  4 ++++
 requirements.txt |  1 +
 setup.py         |  5 +++++
 4 files changed, 22 insertions(+)

diff --git a/README.md b/README.md
index f71bf61..7c127c2 100644
--- a/README.md
+++ b/README.md
@@ -49,3 +49,15 @@ Additionally, you can use dot notation to override specific parameters:
 ```bash
 minppo train configs/stompy_pro.yaml 'rl.num_env_steps=1000'
 ```
+
+You can visualize the environment with:
+
+```bash
+minppo env stompy_pro
+```
+
+Finally, you can run inference on a trained model with:
+
+```bash
+minppo infer stompy_pro 'inference.model_path=path/to/model.pkl'
+```
diff --git a/minppo/cli.py b/minppo/cli.py
index 48d058e..fe72ae9 100644
--- a/minppo/cli.py
+++ b/minppo/cli.py
@@ -2,12 +2,16 @@
 
 import argparse
 
+import colorlogging
+
 from minppo.env import main as env_main
 from minppo.infer import main as infer_main
 from minppo.train import main as train_main
 
 
 def main() -> None:
+    colorlogging.configure()
+
     parser = argparse.ArgumentParser(description="MinPPO CLI")
     parser.add_argument("command", choices=["train", "env", "infer"], help="Command to run")
     args, other_args = parser.parse_known_args()
diff --git a/requirements.txt b/requirements.txt
index cf62032..6eb02e6 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,6 +2,7 @@
 
 # Training
 tqdm
+colorlogging
 
 # Environment
 brax
diff --git a/setup.py b/setup.py
index ebd6cc8..b6607dc 100644
--- a/setup.py
+++ b/setup.py
@@ -41,4 +41,9 @@
     install_requires=requirements,
     tests_require=requirements_dev,
     extras_require={"dev": requirements_dev},
+    entry_points={
+        "console_scripts": [
+            "minppo=minppo.cli:main",
+        ],
+    },
 )