feat(reinforcement learning): ability to train and use an agent trained by RL

jesse/libs/plotly_utils/canvas.py

@@ -0,0 +1,118 @@
+from typing import Literal
+
+import numpy as np
+import plotly.graph_objects as go
+
+from plotly.subplots import make_subplots
+
+import io
+import plotly.io as pio
+from PIL import Image
+
+# pio.kaleido.scope.mathjax = None
+
+
+COLOR_FORMAT = Literal["grayscale", "rgb", "rgba"]
+
+
+class Canvas:
+
+    def __init__(self, width: int = 128, height: int = 128):
+        self.fig = make_subplots(rows=1, cols=1, shared_xaxes=True)
+        self.fig.update_layout(showlegend=False)
+        self.fig.update_layout(
+            margin=dict(l=0, r=0, t=0, b=0),
+            autosize=False,
+            width=width,
+            height=height,
+        )
+        self.fig.update_xaxes(showticklabels=False)
+        self.fig.update_yaxes(showticklabels=False)
+
+    def add_candles(self, candles: np.ndarray, count: int = -1) -> None:
+        self.fig.add_trace(
+            go.Candlestick(
+                open=candles[-count:, 1],
+                close=candles[-count:, 2],
+                high=candles[-count:, 3],
+                low=candles[-count:, 4],
+                increasing={'line': {'color': 'green', 'width': 1}, 'fillcolor': 'green'},
+                decreasing={'line': {'color': 'red', 'width': 1}, 'fillcolor': 'red'},
+            )
+        )
+        self.fig.update(layout_xaxis_rangeslider_visible=False)
+
+    def add_line(self, values: np.ndarray, count: int, **kwargs) -> None:
+        self.fig.add_trace(go.Scatter(y=values[-count:], **kwargs))
+
+    def add_area(
+        self,
+        higher_line: np.ndarray,
+        lower_line: np.ndarray,
+        count: int,
+        fillcolor: str,
+        **kwargs,
+    ) -> None:
+
+        self.fig.add_trace(
+            go.Scatter(
+                y=higher_line[-count:],
+                **kwargs,
+            )
+        )
+        self.fig.add_trace(
+            go.Scatter(
+                y=lower_line[-count:],
+                fill="tonexty",
+                fillcolor=fillcolor,
+                **kwargs,
+            )
+        )
+
+    def to_array(self, format: COLOR_FORMAT = "grayscale") -> np.ndarray:
+        buf = io.BytesIO()
+        self.fig.write_image(
+            buf,
+            engine="kaleido",
+            format="png",
+        )
+        img = Image.open(buf)
+        rgba = np.asarray(img)
+        if format == "rgba":
+            return rgba
+        rgb = rgba[..., :3]
+        alpha = rgba[..., 3]
+        background_color = np.array([255, 255, 255])  # White background
+        alpha = alpha[:, :, np.newaxis] / 255.0
+        rgb_res = rgb * alpha + background_color * (1 - alpha)
+        rgb_res = np.clip(rgb_res, 0, 255).astype(np.uint8)
+        if format == "rgb":
+            return rgb_res
+
+        red, green, blue = (
+            rgb_res[:, :, 0],
+            rgb_res[:, :, 1],
+            rgb_res[:, :, 2],
+        )
+
+        # Standard weights for converting to grayscale
+        gray_values = 0.2989 * red + 0.5870 * green + 0.1140 * blue
+        return gray_values
+
+    def show(self) -> None:
+        buf = io.BytesIO()
+        self.fig.write_image(buf)
+        img = Image.open(buf)
+        img.show()
+
+    def disable_x_axis(self) -> None:
+        self.fig.update_xaxes(showticklabels=False)
+
+    def enable_x_axis(self) -> None:
+        self.fig.update_xaxes(showticklabels=True)
+
+    def disable_y_axis(self) -> None:
+        self.fig.update_yaxes(showticklabels=False)
+
+    def enable_y_axis(self) -> None:
+        self.fig.update_yaxes(showticklabels=True)

jesse/modes/backtest_mode.py

@@ -2,6 +2,8 @@
 from typing import Dict, List, Tuple
 import numpy as np
 import pandas as pd
+from agilerl.algorithms.ppo import PPO
+
 import jesse.helpers as jh
 import jesse.services.metrics as stats
 import jesse.services.selectors as selectors
@@ -215,6 +217,7 @@ def _step_simulator(
         generate_equity_curve: bool = False,
         generate_hyperparameters: bool = False,
         generate_logs: bool = False,
+        agent: PPO | None = None,
 ) -> dict:
     # In case generating logs is specifically demanded, the debug mode must be enabled.
     if generate_logs:
@@ -227,7 +230,7 @@ def _step_simulator(
 
     length = simulation_minutes_length(candles)
     prepare_times_before_simulation(candles)
-    prepare_routes(hyperparameters)
+    prepare_routes(hyperparameters, agent)
 
     # add initial balance
     save_daily_portfolio_balance()
@@ -342,7 +345,7 @@ def prepare_times_before_simulation(candles: dict) -> None:
     store.app.time = first_candles_set[0][0]
 
 
-def prepare_routes(hyperparameters: dict = None) -> None:
+def prepare_routes(hyperparameters: dict = None, agent: PPO | None = None) -> None:
     # initiate strategies
     for r in router.routes:
         # if the r.strategy is str read it from file
@@ -382,6 +385,10 @@ def prepare_routes(hyperparameters: dict = None) -> None:
         # it also injects hyperparameters into self.hp in case the route does not uses any DNAs
         r.strategy._init_objects()
 
+        # override agent in case it passed in the simulation
+        if agent is not None:
+            r.strategy.agent = agent
+
         selectors.get_position(r.exchange, r.symbol).strategy = r.strategy
 
 
@@ -552,6 +559,7 @@ def _skip_simulator(
         generate_equity_curve: bool = False,
         generate_hyperparameters: bool = False,
         generate_logs: bool = False,
+        agent: PPO | None = None,
 ) -> dict:
     # In case generating logs is specifically demanded, the debug mode must be enabled.
     if generate_logs:
@@ -561,7 +569,7 @@ def _skip_simulator(
 
     length = simulation_minutes_length(candles)
     prepare_times_before_simulation(candles)
-    prepare_routes(hyperparameters)
+    prepare_routes(hyperparameters, agent)
 
     # add initial balance
     save_daily_portfolio_balance()

jesse/research/backtest.py

@@ -1,23 +1,26 @@
 from typing import List, Dict
 import copy
 
+from agilerl.algorithms.ppo import PPO
+
 
 def backtest(
-        config: dict,
-        routes: List[Dict[str, str]],
-        extra_routes: List[Dict[str, str]],
-        candles: dict,
-        warmup_candles: dict = None,
-        generate_charts: bool = False,
-        generate_tradingview: bool = False,
-        generate_quantstats: bool = False,
-        generate_hyperparameters: bool = False,
-        generate_equity_curve: bool = False,
-        generate_csv: bool = False,
-        generate_json: bool = False,
-        generate_logs: bool = False,
-        hyperparameters: dict = None,
-        fast_mode: bool = False
+    config: dict,
+    routes: List[Dict[str, str]],
+    extra_routes: List[Dict[str, str]],
+    candles: dict,
+    warmup_candles: dict = None,
+    generate_charts: bool = False,
+    generate_tradingview: bool = False,
+    generate_quantstats: bool = False,
+    generate_hyperparameters: bool = False,
+    generate_equity_curve: bool = False,
+    generate_csv: bool = False,
+    generate_json: bool = False,
+    generate_logs: bool = False,
+    hyperparameters: dict = None,
+    fast_mode: bool = False,
+    agent: PPO | None = None,
 ) -> dict:
     """
     An isolated backtest() function which is perfect for using in research, and AI training
@@ -67,26 +70,28 @@ def backtest(
         generate_hyperparameters=generate_hyperparameters,
         generate_logs=generate_logs,
         fast_mode=fast_mode,
+        agent=agent,
     )
 
 
 def _isolated_backtest(
-        config: dict,
-        routes: List[Dict[str, str]],
-        extra_routes: List[Dict[str, str]],
-        candles: dict,
-        warmup_candles: dict = None,
-        run_silently: bool = True,
-        hyperparameters: dict = None,
-        generate_charts: bool = False,
-        generate_tradingview: bool = False,
-        generate_quantstats: bool = False,
-        generate_csv: bool = False,
-        generate_json: bool = False,
-        generate_equity_curve: bool = False,
-        generate_hyperparameters: bool = False,
-        generate_logs: bool = False,
-        fast_mode: bool = False,
+    config: dict,
+    routes: List[Dict[str, str]],
+    extra_routes: List[Dict[str, str]],
+    candles: dict,
+    warmup_candles: dict = None,
+    run_silently: bool = True,
+    hyperparameters: dict = None,
+    generate_charts: bool = False,
+    generate_tradingview: bool = False,
+    generate_quantstats: bool = False,
+    generate_csv: bool = False,
+    generate_json: bool = False,
+    generate_equity_curve: bool = False,
+    generate_hyperparameters: bool = False,
+    generate_logs: bool = False,
+    fast_mode: bool = False,
+    agent: PPO | None = None,
 ) -> dict:
     from jesse.services.validators import validate_routes
     from jesse.modes.backtest_mode import simulator
@@ -153,6 +158,7 @@ def _isolated_backtest(
         generate_hyperparameters=generate_hyperparameters,
         generate_logs=generate_logs,
         fast_mode=fast_mode,
+        agent=agent,
     )
 
     result = {
@@ -181,6 +187,9 @@ def _isolated_backtest(
         result['hyperparameters'] = backtest_result['hyperparameters']
     if generate_logs:
         result['logs'] = backtest_result['logs']
+    if agent is not None:
+        result['scores'] = store.reinforce_learning.scores()
+        result['experience'] = store.reinforce_learning.experience()
 
     # reset store and config so rerunning would be flawlessly possible
     reset_config()