more improvements

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'
+```

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()

requirements.txt

@@ -2,6 +2,7 @@
 
 # Training
 tqdm
+colorlogging
 
 # Environment
 brax

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",
+        ],
+    },
 )