├── __init__.py
├── gymnax_exchange
    ├── jaxob
    │   ├── __init__.py
    │   ├── README.md
    │   ├── JaxOrderBookWrapper.py
    │   ├── jaxob_constants.py
    │   ├── jaxob_config.py
    │   └── jorderbook.py
    ├── utils
    │   ├── __init__.py
    │   └── utils.py
    ├── jaxlobster
    │   ├── __init__.py
    │   ├── constants.py
    │   └── data_loading.py
    ├── __init__.py
    ├── jaxen
    │   ├── __init__.py
    │   ├── from_JAXMARL
    │   │   ├── __pycache__
    │   │   │   └── multi_agent_env.cpython-312.pyc
    │   │   └── multi_agent_env.py
    │   ├── StatesandParams.py
    │   ├── Speed_test.py
    │   └── base_env.py
    └── jaxrl
    │   ├── __init__.py
    │   └── MARL
    │       ├── config
    │           ├── ippo_rnn_JAXMARL.yaml
    │           ├── ippo_rnn_JAXMARL_simplest_case.yaml
    │           ├── ippo_rnn_JAXMARL_fine_sweep.yaml
    │           ├── ippo_rnn_JAXMARL_2player.yaml
    │           └── PMAP_ippo_rnn_JAXMARL_2player.yaml
    │       └── baseline_eval
    │           ├── config
    │               └── base_config.yaml
    │           ├── simple_triangular_plot.py
    │           ├── plotting_episodes.py
    │           └── plotting_combinations.py
├── Makefile
├── Dockerfile
├── .gitignore
├── README.md
└── requirements.txt


/__init__.py:
--------------------------------------------------------------------------------
1 | 
2 | 


--------------------------------------------------------------------------------
/gymnax_exchange/jaxob/__init__.py:
--------------------------------------------------------------------------------
1 | 


--------------------------------------------------------------------------------
/gymnax_exchange/utils/__init__.py:
--------------------------------------------------------------------------------
1 | 


--------------------------------------------------------------------------------
/gymnax_exchange/jaxlobster/__init__.py:
--------------------------------------------------------------------------------
1 | 


--------------------------------------------------------------------------------
/gymnax_exchange/jaxob/README.md:
--------------------------------------------------------------------------------
1 | 
2 | 


--------------------------------------------------------------------------------
/gymnax_exchange/__init__.py:
--------------------------------------------------------------------------------
1 | try:
2 |   import rich.traceback
3 |   rich.traceback.install()
4 | except ImportError:
5 |   pass


--------------------------------------------------------------------------------
/gymnax_exchange/jaxen/__init__.py:
--------------------------------------------------------------------------------
1 | try:
2 |   import rich.traceback
3 |   rich.traceback.install()
4 | except ImportError:
5 |   pass


--------------------------------------------------------------------------------
/gymnax_exchange/jaxrl/__init__.py:
--------------------------------------------------------------------------------
1 | try:
2 |   import rich.traceback
3 |   rich.traceback.install()
4 | except ImportError:
5 |   pass
6 | 


--------------------------------------------------------------------------------
/gymnax_exchange/jaxen/from_JAXMARL/__pycache__/multi_agent_env.cpython-312.pyc:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/vmohl/JaxMARL-HFT/HEAD/gymnax_exchange/jaxen/from_JAXMARL/__pycache__/multi_agent_env.cpython-312.pyc


--------------------------------------------------------------------------------
/gymnax_exchange/jaxob/JaxOrderBookWrapper.py:
--------------------------------------------------------------------------------
 1 | # ===================================== #
 2 | # ******* Config your own func ******** #
 3 | # ===================================== #
 4 | 
 5 | from gymnax_exchange.jaxob.JaxOrderBookArrays import *
 6 | @partial(jax.jit)
 7 | def get_best_bid(asks, bids):
 8 |     L2_state = get_L2_state(1, asks, bids)
 9 |     return L2_state[0]
10 | 
11 | @partial(jax.jit)
12 | def get_best_bid(asks, bids):
13 |     L2_state = get_L2_state(1, asks, bids)
14 |     return L2_state[1]
15 | 


--------------------------------------------------------------------------------
/gymnax_exchange/jaxrl/MARL/config/ippo_rnn_JAXMARL.yaml:
--------------------------------------------------------------------------------
 1 | # Those hparams which are in a list 
 2 | # apply differently to each ntwrk
 3 | # Should match with NUM_AGENTS_PER_TYPE
 4 | 
 5 | "LR": [0.0004, 0.0004]
 6 | "NUM_ENVS": 512
 7 | "NUM_STEPS": 256
 8 | "GRU_HIDDEN_DIM": 128
 9 | "FC_DIM_SIZE": 128
10 | "TOTAL_TIMESTEPS": 5e7
11 | "UPDATE_EPOCHS": 4
12 | "NUM_MINIBATCHES": 16
13 | "GAMMA": [0.999, 0.999]
14 | "GAE_LAMBDA": [0.99, 0.99]
15 | "CLIP_EPS": 0.2
16 | "SCALE_CLIP_EPS": [False]
17 | "ENT_COEF": [0.001, 0.01]
18 | "VF_COEF": [0.000000000000000000000001, 0.00001]
19 | "MAX_GRAD_NORM": [0.5, 0.5]
20 | "ACTIVATION": ["relu", "relu"]
21 | "ENV_NAME": "JAXMARL-Lob-v0"
22 | "NUM_AGENTS_PER_TYPE": [1, 1]
23 | "SEED": 1
24 | # "ENV_KWARGS": 
25 | #     "LIST_OF_AGENTS_CONFIGS: []
26 | #     "NUM_AGENTS_PER_TYPE": True
27 | # #   "attack_mode": "closest"
28 | "ANNEAL_LR": [True, True]
29 | 
30 | # WandB Params
31 | "ENTITY": "oxford-lob"
32 | "PROJECT": "JAXMARL-LOB_GOOG_2019_two_agents"
33 | "WANDB_MODE" : "online"


--------------------------------------------------------------------------------
/gymnax_exchange/jaxrl/MARL/config/ippo_rnn_JAXMARL_simplest_case.yaml:
--------------------------------------------------------------------------------
 1 | # Those hparams which are in a list 
 2 | # apply differently to each ntwrk
 3 | # Should match with NUM_AGENTS_PER_TYPE
 4 | 
 5 | "LR": [5e-5]
 6 | "NUM_ENVS": 512
 7 | "NUM_STEPS": 256
 8 | "NUM_STEPS_EVAL" : 64
 9 | "GRU_HIDDEN_DIM": 128
10 | "FC_DIM_SIZE": 128
11 | "TOTAL_TIMESTEPS": 1e7
12 | "UPDATE_EPOCHS": 4
13 | "NUM_MINIBATCHES": 16
14 | "GAMMA": [0.999]
15 | "GAE_LAMBDA": [0.99]
16 | "CLIP_EPS": 0.2
17 | "SCALE_CLIP_EPS": [False]
18 | "ENT_COEF": [0.01]
19 | "VF_COEF": [0.001]
20 | "MAX_GRAD_NORM": [0.5]
21 | "ACTIVATION": ["relu"]
22 | "ENV_NAME": "JAXMARL-Lob-v0"
23 | "NUM_AGENTS_PER_TYPE": [1]
24 | "SEED": 0
25 | # "ENV_KWARGS": 
26 | #     "LIST_OF_AGENTS_CONFIGS: []
27 | #     "NUM_AGENTS_PER_TYPE": True
28 | # #   "attack_mode": "closest"
29 | "ANNEAL_LR": [True]
30 | 
31 | # WandB Params
32 | "ENTITY": "oxford-lob"
33 | "PROJECT": "NEWTEST"
34 | "WANDB_MODE" : "online"
35 | 
36 | # High-Level Params
37 | "CALC_EVAL" : True
38 | "CALC_BASELINE" : True


--------------------------------------------------------------------------------
/gymnax_exchange/jaxrl/MARL/config/ippo_rnn_JAXMARL_fine_sweep.yaml:
--------------------------------------------------------------------------------
 1 | # Those hparams which are in a list 
 2 | # apply differently to each ntwrk
 3 | # Should match with NUM_AGENTS_PER_TYPE
 4 | 
 5 | "LR": [5e-5]
 6 | "NUM_ENVS": 512
 7 | "NUM_STEPS": 512
 8 | "GRU_HIDDEN_DIM": 128
 9 | "FC_DIM_SIZE": 128
10 | "TOTAL_TIMESTEPS": 1e7
11 | "UPDATE_EPOCHS": 4
12 | "NUM_MINIBATCHES": 16
13 | "GAMMA": [0.999]
14 | "GAE_LAMBDA": [0.99]
15 | "CLIP_EPS": 0.2
16 | "SCALE_CLIP_EPS": [False]
17 | "ENT_COEF": [0.01]
18 | "VF_COEF": [0.001]
19 | "MAX_GRAD_NORM": [0.5]
20 | "ACTIVATION": ["relu"]
21 | "ENV_NAME": "JAXMARL-Lob-v0"
22 | "NUM_AGENTS_PER_TYPE": [1]
23 | "SEED": 0
24 | # "ENV_KWARGS": 
25 | #     "LIST_OF_AGENTS_CONFIGS: []
26 | #     "NUM_AGENTS_PER_TYPE": True
27 | # #   "attack_mode": "closest"
28 | "ANNEAL_LR": [True]
29 | 
30 | # WandB Params
31 | "ENTITY": "oxford-lob"
32 | "PROJECT": "DebugConfig"
33 | "WANDB_MODE" : "online"
34 | 
35 | # High-Level Params
36 | "CALC_EVAL" : True
37 | "Calculate Baseline" : True
38 | 
39 | 
40 | 
41 | 


--------------------------------------------------------------------------------
/gymnax_exchange/jaxrl/MARL/config/ippo_rnn_JAXMARL_2player.yaml:
--------------------------------------------------------------------------------
 1 | # Those hparams which are in a list 
 2 | # apply differently to each ntwrk
 3 | # Should match with NUM_AGENTS_PER_TYPE
 4 | 
 5 | "LR": [0.0004,0.0004]
 6 | "NUM_ENVS": 4096 #128
 7 | "NUM_STEPS": 128
 8 | "NUM_STEPS_EVAL" : 64
 9 | "GRU_HIDDEN_DIM": 128
10 | "FC_DIM_SIZE": 128
11 | "TOTAL_TIMESTEPS": 5e7
12 | "UPDATE_EPOCHS": 4
13 | "NUM_MINIBATCHES": 16
14 | "GAMMA": [0.999,0.999]
15 | "GAE_LAMBDA": [0.99,0.99]
16 | "CLIP_EPS": 0.2
17 | "SCALE_CLIP_EPS": [False,False]
18 | "ENT_COEF": [0.001,0.001]
19 | "VF_COEF": [1e-11,1e-11]
20 | "MAX_GRAD_NORM": [0.5,0.5]
21 | "ACTIVATION": ["relu","relu"]
22 | "ENV_NAME": "JAXMARL-Lob-v0"
23 | "NUM_AGENTS_PER_TYPE": [1,1]
24 | "SEED": 2
25 | "N_DEVICES": 1
26 | # "ENV_KWARGS": 
27 | #     "LIST_OF_AGENTS_CONFIGS: []
28 | #     "NUM_AGENTS_PER_TYPE": True
29 | # #   "attack_mode": "closest"
30 | "ANNEAL_LR": [True,True]
31 | 
32 | # WandB Params
33 | "ENTITY": "oxford-lob"
34 | "PROJECT": "test2"
35 | "WANDB_MODE" : "online"
36 | 
37 | # High-Level Params
38 | "CALC_EVAL" : False
39 | "EvalTimePeriod" : "2024_Eval"
40 | "Calculate Baseline" : False
41 | "Timing" : False
42 | 
43 | 
44 | 
45 | 


--------------------------------------------------------------------------------
/gymnax_exchange/jaxrl/MARL/config/PMAP_ippo_rnn_JAXMARL_2player.yaml:
--------------------------------------------------------------------------------
 1 | # Those hparams which are in a list 
 2 | # apply differently to each ntwrk
 3 | # Should match with NUM_AGENTS_PER_TYPE
 4 | 
 5 | "LR": [0.0004,0.0004]
 6 | "NUM_ENVS": 32768 #32768 #128
 7 | "NUM_STEPS": 64
 8 | "NUM_STEPS_EVAL" : 64
 9 | "GRU_HIDDEN_DIM": 128
10 | "FC_DIM_SIZE": 128
11 | "TOTAL_TIMESTEPS": 1e10
12 | "UPDATE_EPOCHS": 4
13 | "NUM_MINIBATCHES": 16
14 | "GAMMA": [0.999,0.999]
15 | "GAE_LAMBDA": [0.99,0.99]
16 | "CLIP_EPS": 0.2
17 | "SCALE_CLIP_EPS": [False,False]
18 | "ENT_COEF": [0.001,0.001]
19 | "VF_COEF": [1e-11,1e-11]
20 | "MAX_GRAD_NORM": [0.5,0.5]
21 | "ACTIVATION": ["relu","relu"]
22 | "ENV_NAME": "JAXMARL-Lob-v0"
23 | "NUM_AGENTS_PER_TYPE": [10,10]
24 | "SEED": 2
25 | "N_DEVICES": 8
26 | # "ENV_KWARGS": 
27 | #     "LIST_OF_AGENTS_CONFIGS: []
28 | #     "NUM_AGENTS_PER_TYPE": True
29 | # #   "attack_mode": "closest"
30 | "ANNEAL_LR": [True,True]
31 | 
32 | # WandB Params
33 | "ENTITY": "oxford-lob"
34 | "PROJECT": "PMAP_SPEED" #"DEBUG"
35 | "WANDB_MODE" : "online"
36 | 
37 | # High-Level Params
38 | "CALC_EVAL" : False
39 | "EvalTimePeriod" : "2024_Eval"
40 | "Calculate Baseline" : False
41 | "Timing" : True
42 | 
43 | 
44 | 
45 | 


--------------------------------------------------------------------------------
/Makefile:
--------------------------------------------------------------------------------
 1 | NVCC_RESULT := $(shell which nvcc 2> NULL; rm NULL)
 2 | NVCC_TEST := $(notdir $(NVCC_RESULT))
 3 | ifeq ($(NVCC_TEST),nvcc)
 4 | GPUS=--gpus '"device=7"'
 5 | else
 6 | GPUS=
 7 | endif
 8 | 
 9 | 
10 | # Set flag for docker run command
11 | MYUSER=myuser
12 | SERVER_NAME = $(shell hostname)
13 | # If using flair12 server, set data directory to /homes/80/sascha/data, otherwise assume data is on same level as the repo
14 | ifeq ($(SERVER_NAME),flair-node-12)
15 | DATADIR=/homes/80/sascha/data
16 | else
17 | DATADIR=~/data
18 | endif
19 | BASE_FLAGS=-it --rm -v ${PWD}:/home/$(MYUSER) -v $(DATADIR):/home/$(MYUSER)/data --shm-size 20G 
20 | PORT_FLAGS= -p 8074:80 -p 8075:6006
21 | RUN_FLAGS=$(GPUS) $(BASE_FLAGS) $(PORT_FLAGS)
22 | BASIC_FLAGS=$(GPUS) $(BASE_FLAGS)
23 | 
24 | 
25 | DOCKER_IMAGE_NAME = jaxmarl_lob
26 | IMAGE = $(DOCKER_IMAGE_NAME):latest #  for working image: IMAGE = $(DOCKER_IMAGE_NAME):working or 
27 | DOCKER_RUN=docker run $(RUN_FLAGS) $(IMAGE)
28 | DOCKER_RUN_BASIC=docker run --gpus "device=$(gpu)" $(BASE_FLAGS) $(IMAGE)
29 | USE_CUDA = $(if $(GPUS),true,false)
30 | ID = $(shell id -u)
31 | 
32 | # make file commands
33 | build:
34 | 	DOCKER_BUILDKIT=1 docker build -f Dockerfile  --build-arg USE_CUDA=$(USE_CUDA) --build-arg MYUSER=$(MYUSER) --build-arg UID=$(ID) --tag $(IMAGE) --progress=plain ${PWD}/. 
35 | 	
36 | # --no-cache
37 | 	
38 | 
39 | run:
40 | 	$(DOCKER_RUN) /bin/bash
41 | 
42 | test:
43 | 	$(DOCKER_RUN) /bin/bash -c "pytest ./tests/"
44 | 
45 | ppo:
46 | 	$(DOCKER_RUN_BASIC) /bin/bash -c "python3 ./gymnax_exchange/jaxrl/MARL/ippo_rnn_JAXMARL.py"
47 | 
48 | workflow-test:
49 | 	# without -it flag
50 | 	docker run --rm -v ${PWD}:/home/workdir --shm-size 20G $(IMAGE) /bin/bash -c "pytest ./tests/"


--------------------------------------------------------------------------------
/gymnax_exchange/jaxob/jaxob_constants.py:
--------------------------------------------------------------------------------
 1 | 
 2 | from enum import Enum
 3 | 
 4 | class MaxInt(Enum):
 5 |     _64_Bit_Signed=2_147_483_647
 6 |     _32_Bit_Signed=32_767
 7 | 
 8 | # Actual constants, will never change.
 9 | INITID=-2
10 | DUMMYID=-888888
11 | EMPTY_SLOT=-1
12 | 
13 | ORDERBOOK_FEAT=6
14 | TRADE_FEAT=8
15 | NS_PER_SEC=1e9
16 | 
17 | # Default values for the config.
18 | 
19 | NTRADE_CAP=100
20 | NORDER_CAP=100
21 | STARTOFDAY=[34200 , 0]
22 | ENDOFDAY=[57600 , 0]
23 | TEST_TIME=[44444,44444]
24 | 
25 | 
26 | # LOBSTER Message types, 
27 | class MessageType(Enum):
28 |     LIMIT=1  
29 |     CANCEL=2 
30 |     DELETE=3 
31 |     MATCH=4
32 |     HIDDEN=5
33 |     AUCTION=6
34 |     HALT=7
35 | 
36 | class OrderSideFeat(Enum):
37 |     P=0 #Price
38 |     Q=1 #Quantity
39 |     OID=2 #Order ID
40 |     TID=3 # Trade ID
41 |     SEC=4 #Seconds
42 |     NSEC=5 #Nanoseconds 
43 | 
44 | class TradesFeat(Enum):
45 |     P=0 #Price
46 |     Q=1 #Quantity
47 |     PASS_OID=2 #Order ID
48 |     AGRS_OID=3 # Trade ID
49 |     SEC=4 #Seconds
50 |     NSEC=5 #Nanoseconds 
51 |     PASS_TID=6 #Order ID
52 |     AGRS_TID=7 # Trade ID
53 | 
54 | 
55 | 
56 | class BidAskSide(Enum):
57 |     BID=1
58 |     ASK=-1
59 | 
60 | 
61 | #Define as static...? Might solve 
62 | class CancelMode(Enum):
63 |     STRICT_BY_ID=0 #Cancel only if ID matches. 
64 |     INCLUDE_INITS=1 #Cancel only init orders if ID does not match
65 |     CANCEL_UNIFORM=2  #Pick a random order at the right price level to cancel 
66 |     CANCEL_UNIFORM_AND_LARGE=3 # Unused for now
67 | 
68 | SEED= 42 # the meaning of life. 
69 | 
70 | #TODO: flag on behaviour of type market either limit or far touch. 
71 | 
72 | class SimulatorMode(Enum):
73 |     GENERAL_EXCHANGE=0
74 |     LOBSTER_INTERPRETER=1
75 | 
76 | class LOBMSGFEAT(Enum):
77 |     Type=0  # Message type
78 |     Side=1  # Order side (buy/sell)
79 |     Quant=2 # Quantity
80 |     Price=3 # Price
81 |     OID=4   # Order ID
82 |     TID=5   # Trade ID
83 |     TS=6    # Timestamp seconds
84 |     TNS=7   # Timestamp nanoseconds


--------------------------------------------------------------------------------
/Dockerfile:
--------------------------------------------------------------------------------
 1 | FROM nvcr.io/nvidia/jax:25.01-py3
 2 | 
 3 | # Create user
 4 | ARG UID
 5 | ARG MYUSER
 6 | RUN useradd -m -u $UID --create-home ${MYUSER} && \
 7 |     echo "${MYUSER}: ${MYUSER}" | chpasswd && \
 8 |     adduser ${MYUSER} sudo && \
 9 |     # mkdir -p /home/${MYUSER}/.local/bin && \
10 |     chown -R  ${MYUSER}:${MYUSER} /home/${MYUSER}
11 |     
12 | # default workdir
13 | WORKDIR /home/${MYUSER}/
14 | COPY --chown=${MYUSER} --chmod=777 . .
15 | 
16 | # install tmux
17 | RUN apt-get update && \
18 |     apt-get install -y tmux && \
19 |     apt-get install -y \
20 |     apt-transport-https \
21 |     ca-certificates \
22 |     gnupg \
23 |     libgl1 \
24 |     p7zip-full \
25 |     unrar \
26 |     htop \
27 |     graphviz \
28 |     libcupti-dev
29 | # Copy requirements.txt and verify its contents
30 | COPY --chown=${MYUSER}:${MYUSER} requirements.txt /home/${MYUSER}/AlphaTrade/
31 | RUN ls -l && cat requirements.txt
32 | 
33 | #jaxmarl from source if needed, all the requirements
34 | # RUN pip install -e .[algs,dev]
35 | RUN pip install --no-cache-dir -r requirements.txt
36 | 
37 | 
38 | RUN wget https://go.dev/dl/go1.24.5.linux-amd64.tar.gz
39 | RUN rm -rf /usr/local/go
40 | RUN tar -C /usr/local -xzf go1.24.5.linux-amd64.tar.gz
41 | 
42 | RUN export PATH=$PATH:/usr/local/go/bin
43 | 
44 | USER ${MYUSER}
45 | 
46 | #disabling preallocation
47 | RUN export XLA_PYTHON_CLIENT_PREALLOCATE=false
48 | #safety measures
49 | RUN export XLA_PYTHON_CLIENT_MEM_FRACTION=0.25 
50 | RUN export TF_FORCE_GPU_ALLOW_GROWTH=true
51 | 
52 | # Add .local/bin to PATH
53 | RUN echo 'export PATH=$PATH:/home/duser/.local/bin' >> ~/.bashrc
54 | 
55 | 
56 | # Ensure home directory is on the Python path
57 | ENV PYTHONPATH="/home/${MYUSER}:$PYTHONPATH"
58 | RUN export PATH="$HOME/.local/bin:$PATH"
59 | RUN export PATH="$HOME/.local/lib:$PATH"
60 | # Uncomment below if you want jupyter 
61 | # RUN pip install jupyterlab
62 | 
63 | #for secrets and debug
64 | ENV WANDB_API_KEY=""
65 | ENV WANDB_ENTITY=""
66 | RUN git config --global --add safe.directory /home/${MYUSER}
67 | 
68 | 


--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
  1 | **.csv
  2 | .idea/**/
  3 | **.pyc
  4 | .idea/**
  5 | .spyproject/**
  6 | .DS_Store
  7 | **.log
  8 | /train/output/**
  9 | **.pkl
 10 | /train/wandb/**
 11 | **/wandb/**
 12 | /gym_exchange/outputs/**
 13 | .bash_history
 14 | jax_rwkv/**
 15 | .bash_history
 16 | 
 17 | gym_exchange/.DS_Store
 18 | gym_exchange/data_orderbook_adapter/.DS_Store
 19 | gym_exchange/orderbook/.DS_Store
 20 | gymnax_exchange/.DS_Store
 21 | **.xml
 22 | **/.idea/**/
 23 | gymnax_exchange/.idea/
 24 | gymnax_exchange/.idea/**
 25 | **/.DS_Store
 26 | gymnax_exchange/.DS_Store
 27 | .vscode/**
 28 | **/.ipynb_checkpoints/**
 29 | gymnax_exchange/jaxob/.ipynb_checkpoints/**
 30 | Miniconda3-latest-Linux-x86_64.sh
 31 | nohup.out
 32 | program.prof
 33 | **/gymnax/**
 34 | params_file
 35 | params_file_2023-07-08_15-22-20
 36 | params_file_2023-07-08_19-29-37
 37 | re_1.py
 38 | re_2.ipynb
 39 | **.ckpt
 40 | params_file**
 41 | note.md
 42 | results_file**
 43 | 
 44 | *.png
 45 | *.npy
 46 | *.nfs*
 47 | 
 48 | Dockerfile
 49 | Dockerfile_twogymnax_exchange/jaxrl/pre_trained_weights/*.model
 50 | tensorboard_logs/
 51 | tensorboard_logs/**
 52 | 
 53 | old_version_tb_logs/
 54 | old_version_tb_logs/**
 55 | report*
 56 | .nv/**
 57 | .nv/ */
 58 | 
 59 | 
 60 | eval_results/**
 61 | 
 62 | .cache/
 63 | .cache/**
 64 | .keras/
 65 | .keras/**
 66 | .local/
 67 | .local/**
 68 | 
 69 | .vscode-server/**
 70 | .outputs/**
 71 | .mycache/**
 72 | .ssh/**
 73 | .chache/**
 74 | .outputs/**
 75 | /training_oneDay/
 76 | /training_oneDay/**
 77 | /outputs/**
 78 | /.dotnet/**
 79 | .dotnet/
 80 | .dotnet/**
 81 | .vscode/
 82 | .vscode-server/**
 83 | .gnupg/
 84 | .gnupg/**
 85 | outputs/
 86 | outputs/**
 87 | 
 88 | .ssh/
 89 | .ssh/**
 90 | 
 91 | .venv/
 92 | .venv/**
 93 | /mycache/**
 94 | /cache/**
 95 | 
 96 | /gymnax_exchange/jaxen/Timing_speed/** */
 97 | /gymnax_exchange/jaxen/Timing_speed/**
 98 | .netrc
 99 | /gymnax_exchange/jaxen/output/**
100 | /old/.netrc
101 | 
102 | /go/*
103 | .config/*
104 | saved_npz/**
105 | saved_npz/
106 | .netrc
107 | /checkpoints/**
108 | .gitconfig
109 | 
110 | test.py


--------------------------------------------------------------------------------
/gymnax_exchange/jaxlobster/constants.py:
--------------------------------------------------------------------------------
 1 | from dataclasses import dataclass
 2 | import numpy as np
 3 | 
 4 | TIME_COL = "<time>"
 5 | EVENT_TYPE_COL = "<event_type>"
 6 | ORDER_ID_COL = "<order_id>"
 7 | SIZE_COL = "<size>"
 8 | PRICE_COL = "<price>"
 9 | DIRECTION_COL = "<direction>"
10 | 
11 | MESSAGE_TOKEN_DTYPE_MAP = {
12 |     TIME_COL: int,  
13 |     EVENT_TYPE_COL: int,
14 |     ORDER_ID_COL: int,
15 |     SIZE_COL: int,
16 |     PRICE_COL: int,
17 |     DIRECTION_COL: int
18 | }
19 | MESSAGE_TOKEN_TYPES = list(MESSAGE_TOKEN_DTYPE_MAP.keys())
20 | 
21 | def get_orderbook_token_types(levels: int) -> list[str]: 
22 |     return np.array([
23 |                 [f"<ask_price_{i}>", f"<ask_size_{i}>", f"<bid_price_{i}>", f"<bid_size_{i}>"]
24 |                 for i in range(1, levels + 1)]
25 |             ).flatten().tolist()
26 | 
27 | @dataclass
28 | class MambaTrainArgs:
29 |     train_data_dir: str = "./data/raw/"
30 |     eval_data_dir: str = "./data/test/"
31 |     file_filter_train: str = ""
32 |     file_filter_eval: str = ""
33 |     save_path: str = "./models/mamba2"
34 |     nmsgs: int = 50
35 |     only_use_message_orderbook_matches: bool = True
36 | 
37 |     tokenizer_file: str = "tokenizers/lob_tok_with_time_diff.json"
38 | 
39 |     wandb_online: bool = True
40 |     wandb_project: str = "lobgen"
41 |     wandb_entity: str = "gereon-franken-oxford"
42 | 
43 | 
44 | 
45 | @dataclass
46 | class MambaInferenceArgs:
47 |     model_path: str
48 |     tokenizer_path: str = "tokenizers/lob_tok_messages.json"
49 |     is_sharded: bool = False
50 |     test_dir: str = "data/GOOG/test/"
51 |     test_filter: str = "2018-12-31"
52 |     genlen: int = 100
53 |     iterations: int = 10
54 |     temperature: float = 1.0
55 |     topk: int = 50
56 |     topp: float = 1.0
57 |     minp: float = 0.0
58 |     repetition_penalty: float = 1.0
59 |     batch: int = 1
60 | 
61 | @dataclass
62 | class MambaBenchmarkingArgs(MambaInferenceArgs):
63 |     data_dir: str = "data/GOOG/2018/"
64 |     data_time_stamp: str = "2018-12-31"
65 |     save_path: str = "gen_data/"
66 | 
67 | @dataclass
68 | class TokenizerTrainArgs:
69 |     data_dir: str = "./data/raw/"
70 |     file_filter: str = "*.csv"
71 |     save_path: str = "./tokenizers/lob_tok.json"
72 |     vocab_size: int = 10_000


--------------------------------------------------------------------------------
/gymnax_exchange/jaxrl/MARL/baseline_eval/config/base_config.yaml:
--------------------------------------------------------------------------------
 1 | # Those hparams which are in a list 
 2 | # apply differently to each ntwrk
 3 | # Should match with NUM_AGENTS_PER_TYPE
 4 | 
 5 | "NUM_ENVS": 4096
 6 | "NUM_STEPS": 128
 7 | "TINY_RUN": False
 8 | 
 9 | # "NUM_ENVS": 8
10 | # "NUM_STEPS": 64
11 | # "TINY_RUN": True
12 | 
13 | 
14 | "ENV_NAME": "JAXMARL-Lob-v0"
15 | "NUM_AGENTS_PER_TYPE": [1,1]
16 | "SEED": 34
17 | "N_DEVICES": 1
18 | "FIXED_ACTIONS" : [[0],[0]]
19 | 
20 | # WandB Params
21 | "ENTITY": "oxford-lob"
22 | "PROJECT": "Baseline_JAXMARL"
23 | "WANDB_MODE" : "disabled"
24 | 
25 | # High-Level Params
26 | "CALC_EVAL" : False
27 | "EvalTimePeriod" : "2024_Eval" #"2024"
28 | 
29 | 
30 | #Network params that must be manually matched with the config used in the original training
31 | "GRU_HIDDEN_DIM": 128
32 | "FC_DIM_SIZE": 128
33 | #Training params which are irellevant for the eval.
34 | "ANNEAL_LR": True
35 | "LR": 20
36 | "MAX_GRAD_NORM" : 25
37 | 
38 | 
39 | "RESTORE_PROJECT" : "v4_2Player_various_act_reward_spaces"
40 | "RESTORE_RUN" : "fluent-sweep-3"
41 | 
42 | 
43 | 
44 | "AGENT_CONFIGS" : {"MarketMaking" :{"inv_penalty": "quadratic",
45 |                                     "skew_multiplier": 10,
46 |                                     "action_space":"fixed_quants", #"spread_skew",,"fixed_quants"simple
47 |                                     "reward_space" : "spooner", # "spooner"buy_sell_pnl
48 |                                   },
49 |                   "Execution" : {"reward_lambda": 0.0,
50 |                                 "fixed_quant_value": 10, #20 on fixed quants
51 |                                 "action_space": "fixed_quants_complex", #fixed_quants,fixed_quants_complex
52 |                                  "task_size": 600,
53 |                                 "doom_price_penalty": 0.1,}
54 |                 }
55 | "BASELINE_CONFIGS" : {"MarketMaking" :{"action_space":"AvSt", #"spread_skew",,"fixed_quants"simple
56 |                                     },
57 |                   "Execution" : {"action_space": "twap"}
58 |                 }
59 | 
60 | 
61 | 
62 | 
63 |       #   """Action function for the fixed Quant Action space
64 |       #   Pick for a ladder of quant execution options
65 |       #   Always send 4 messages
66 |       #   0 = No trade
67 |       #   1=      # FT
68 |       #   2=     # M
69 |       #   3=    # NT
70 |       #   4=    # PP
71 |       #   5=    # FT*2 quant
72 |       #   6=    # M*2 quant
73 |       #   7=    # NT*2 quant
74 |       #   8=    # PP*2 quant
75 |       #   9=    # FT*5 quant
76 |       #   10=   # M*5 quant
77 |       #   11=   # NT*5 quant
78 |       #   12=   # PP*5 quant
79 | 
80 |         
81 |       #  """
82 | 
83 | 
84 |         # Mapping:
85 |         # 0: tight spread, bid skew
86 |         # 1: tight spread, neutral 
87 |         # 2: tight spread, ask skew
88 |         # 3: wide spread, bid skew
89 |         # 4: wide spread, neutral
90 |         # 5: wide spread, ask skew
91 |         # '''


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # JaxMARL-HFT: GPU-Accelerated Multi-Agent Reinforcement Learning for High-Frequency Trading
 2 | 
 3 | A JAX-based framework for multi-agent reinforcement learning for high-frequency trading, based on the [JAX-LOB simulator](https://github.com/KangOxford/jax-lob) and an extension of [JaxMARL](https://github.com/FLAIROx/JaxMARL) to the financial trading domain.
 4 | 
 5 | ## Key Features
 6 | 
 7 | - **GPU-Accelerated**: Built on JAX for high-performance parallel computation with JIT compilation
 8 | - **Two levels of Parallelization**: Parallel processing across episodes and agent types using `vmap`
 9 | - **Multi-Agent RL**: Supports market making, execution, and directional trading agents
10 | - **LOBSTER Data Integration**: Real market data support with efficient GPU memory usage
11 | - **Scalable**: Handles thousands of parallel environments 
12 | - **Heterogeneous Agents**: Supports different observation/action spaces
13 | 
14 | ## Quick Start
15 | 
16 | ### Docker Setup
17 | 
18 | ```bash
19 | # Set up data directory
20 | mkdir -p ~/data
21 | ```
22 | 
23 | **Note**: Configure the Makefile for your specific environment (GPU device, data directory path, etc.)
24 | 
25 | ```bash
26 | # Build and run with Docker
27 | make build
28 | make run
29 | ```
30 | 
31 | ### Training
32 | 
33 | ```bash
34 | # Run IPPO training
35 | python3 gymnax_exchange/jaxrl/MARL/ippo_rnn_JAXMARL.py
36 | ```
37 | 
38 | ## Agent Types
39 | 
40 | ### Market Making Agents
41 | - **Purpose**: Provide liquidity by posting bid/ask orders
42 | - **Action Spaces**: Multiple discrete action spaces (spread_skew, fixed_quants, AvSt, directional_trading, simple)
43 | - **Reward Functions**: Various PnL-based rewards with configurable inventory penalties
44 | 
45 | ### Execution Agents  
46 | - **Purpose**: Execute large orders with minimal market impact
47 | - **Action Spaces**: Discrete quantity selection at reference prices (fixed_quants, fixed_prices, complex variants)
48 | - **Reward Functions**: Slippage-based with configurable end-of-episode penalties
49 | 
50 | ### Directional Trading
51 | - **Purpose**: Simple directional trading strategy
52 | - **Action Spaces**: Bid/ask at best prices or no action
53 | - **Reward Function**: Portfolio value
54 | - **Note:** Uses the same class as the market making agent
55 | 
56 | ## Repository Structure
57 | 
58 | ```
59 | gymnax_exchange/
60 | ├── jaxen/            # Environment implementations
61 | │   ├── marl_env.py  # Multi-agent RL environment
62 | │   ├── mm_env.py    # Market making (and directional trading) environment  
63 | │   └── exec_env.py  # Execution environment
64 | ├── jaxrl/           # Reinforcement learning algorithms
65 | │   └── MARL/        # IPPO implementation
66 | ├── jaxob/           # Order book implementation
67 | └── jaxlobster/      # LOBSTER data integration
68 | ```
69 | 
70 | ## Configuration
71 | 
72 | The framework uses a comprehensive configuration system with dataclasses for different components:
73 | 
74 | ### Core Configuration Classes
75 | 
76 | - **`MultiAgentConfig`**: Main configuration combining world and agent settings
77 | - **`World_EnvironmentConfig`**: Global environment parameters (data paths, episode settings, market hours)
78 | - **`MarketMaking_EnvironmentConfig`**: Market making and directional trading agent configuration (action spaces, reward functions, observation spaces)
79 | - **`Execution_EnvironmentConfig`**: Execution agent configuration (task types, action spaces, reward parameters)
80 | 
81 | ### Training Configuration
82 | 
83 | Edit YAML files in `gymnax_exchange/jaxrl/MARL/config/` to customize:
84 | - Number of parallel environments (default: 4096)
85 | - Training parameters (steps, learning rates, etc.)
86 | - Agent configurations (action spaces, reward functions)
87 | - Market data settings (resolution, episode length)
88 | 
89 | ## Requirements
90 | 
91 | - Python 3.8+
92 | - CUDA-compatible GPU (recommended)
93 | - JAX, Flax, and related dependencies (see `requirements.txt`)
94 | 
95 | 


--------------------------------------------------------------------------------
/gymnax_exchange/jaxen/from_JAXMARL/multi_agent_env.py:
--------------------------------------------------------------------------------
  1 | """ 
  2 | Abstract base class for multi agent gym environments with JAX
  3 | Based on the Gymnax and PettingZoo APIs
  4 | 
  5 | """
  6 | 
  7 | import jax
  8 | import jax.numpy as jnp
  9 | from typing import Dict
 10 | import chex
 11 | from functools import partial
 12 | from flax import struct
 13 | from typing import Tuple, Optional
 14 | jax.config.update('jax_disable_jit', False)
 15 | 
 16 | 
 17 | @struct.dataclass
 18 | class State:
 19 |     done: chex.Array
 20 |     step: int
 21 | 
 22 | 
 23 | class MultiAgentEnv(object):
 24 |     """Jittable abstract base class for all jaxmarl Environments."""
 25 | 
 26 |     def __init__(
 27 |         self,
 28 |         num_agents: int,
 29 |     ) -> None:
 30 |         """
 31 |         num_agents (int): maximum number of agents within the environment, used to set array dimensions
 32 |         """
 33 |         self.num_agents = num_agents
 34 |         self.observation_spaces = dict()
 35 |         self.action_spaces = dict()
 36 | 
 37 |     @partial(jax.jit, static_argnums=(0,))
 38 |     def reset(self, key: chex.PRNGKey) -> Tuple[Dict[str, chex.Array], State]:
 39 |         """Performs resetting of the environment."""
 40 |         raise NotImplementedError
 41 | 
 42 |     @partial(jax.jit, static_argnums=(0,))
 43 |     def step(
 44 |         self,
 45 |         key: chex.PRNGKey,
 46 |         state: State,
 47 |         actions: Dict[str, chex.Array],
 48 |         reset_state: Optional[State] = None,
 49 |     ) -> Tuple[Dict[str, chex.Array], State, Dict[str, float], Dict[str, bool], Dict]:
 50 |         """Performs step transitions in the environment. Resets the environment if done.
 51 |         To control the reset state, pass `reset_state`. Otherwise, the environment will reset randomly."""
 52 | 
 53 |         key, key_reset = jax.random.split(key)
 54 |         obs_st, states_st, rewards, dones, infos = self.step_env(key, state, actions)
 55 | 
 56 |         if reset_state is None:
 57 |             obs_re, states_re = self.reset(key_reset)
 58 |         else:
 59 |             states_re = reset_state
 60 |             obs_re = self.get_obs(states_re)
 61 | 
 62 |         # Auto-reset environment based on termination
 63 |         states = jax.tree.map(
 64 |             lambda x, y: jax.lax.select(dones["__all__"], x, y), states_re, states_st
 65 |         )
 66 |         obs = jax.tree.map(
 67 |             lambda x, y: jax.lax.select(dones["__all__"], x, y), obs_re, obs_st
 68 |         )
 69 |         return obs, states, rewards, dones, infos
 70 | 
 71 |     def step_env(
 72 |         self, key: chex.PRNGKey, state: State, actions: Dict[str, chex.Array]
 73 |     ) -> Tuple[Dict[str, chex.Array], State, Dict[str, float], Dict[str, bool], Dict]:
 74 |         """Environment-specific step transition."""
 75 |         raise NotImplementedError
 76 | 
 77 |     def get_obs(self, state: State) -> Dict[str, chex.Array]:
 78 |         """Applies observation function to state."""
 79 |         raise NotImplementedError
 80 | 
 81 |     def observation_space(self, agent: str):
 82 |         """Observation space for a given agent."""
 83 |         return self.observation_spaces[agent]
 84 | 
 85 |     def action_space(self, agent: str):
 86 |         """Action space for a given agent."""
 87 |         return self.action_spaces[agent]
 88 | 
 89 |     @partial(jax.jit, static_argnums=(0,))
 90 |     def get_avail_actions(self, state: State) -> Dict[str, chex.Array]:
 91 |         """Returns the available actions for each agent."""
 92 |         raise NotImplementedError
 93 | 
 94 |     @property
 95 |     def name(self) -> str:
 96 |         """Environment name."""
 97 |         return type(self).__name__
 98 | 
 99 |     @property
100 |     def agent_classes(self) -> dict:
101 |         """Returns a dictionary with agent classes, used in environments with hetrogenous agents.
102 | 
103 |         Format:
104 |             agent_base_name: [agent_base_name_1, agent_base_name_2, ...]
105 |         """
106 |         raise NotImplementedError


--------------------------------------------------------------------------------
/gymnax_exchange/jaxen/StatesandParams.py:
--------------------------------------------------------------------------------
  1 | """
  2 | State and Parameter Definitions for Multi-Agent Limit Order Book Trading
  3 | 
  4 | University of Oxford
  5 | Corresponding Author: 
  6 | Valentin Mohl (valentin.mohl@cs.ox.ac.uk)
  7 | Reuben Leyland (Reuben.leyland@sky.com)
  8 | Sascha Frey (sascha.frey@st-hughs.ox.ac.uk)
  9 | 
 10 | 
 11 | Module Description
 12 | This module defines all state and parameter classes used across the multi-agent 
 13 | reinforcement learning environment for limit order book trading. It provides 
 14 | structured data classes for managing environment states, agent states, and 
 15 | configuration parameters using JAX-compatible dataclasses.
 16 | 
 17 | Key Components
 18 | LoadedEnvState:    Base state class for the loaded environment, containing
 19 |                    raw order book data, trades, and timing information, which is loaded from the data files.
 20 | WorldState:        Extended state class with market information like best bids/asks,
 21 |                    mid price, and order ID counter etc.
 22 | MultiAgentState:   Combined state class managing the shared world state and
 23 |                    individual agent states for multi-agent coordination.
 24 | LoadedEnvParams:   Base parameters class for environment data and initialization.
 25 | MultiAgentParams:  Combined parameters class for multi-agent coordination.
 26 | MMEnvParams:       Market making and directional trading agent-specific parameters.
 27 | ExecEnvParams:     Execution agent-specific parameters.
 28 | MMEnvState:        Market making and directional trading agent state with inventory and position tracking.
 29 | ExecEnvState:      Execution agent state with task-specific information.
 30 | 
 31 | State Hierarchy
 32 | - LoadedEnvState: Base environment state
 33 |   - WorldState: Extended with market information
 34 |     - MultiAgentState: Combined with agent states
 35 | - MMEnvState: Market making agent state
 36 | - ExecEnvState: Execution agent state
 37 | 
 38 | Parameter Hierarchy  
 39 | - LoadedEnvParams: Base environment parameters
 40 |   - MultiAgentParams: Combined with agent parameters
 41 | - MMEnvParams: Market making agent parameters
 42 | - ExecEnvParams: Execution agent parameters
 43 | """
 44 | 
 45 | import jax.numpy as jnp
 46 | from flax import struct
 47 | from typing import Any
 48 | import chex
 49 | 
 50 | 
 51 | 
 52 | ########################################################################################
 53 | ########################################################################################
 54 | # States
 55 | ########################################################################################
 56 | ########################################################################################
 57 | 
 58 | @struct.dataclass
 59 | class LoadedEnvState:
 60 |     ask_raw_orders: chex.Array
 61 |     bid_raw_orders: chex.Array
 62 |     trades: chex.Array
 63 |     init_time: chex.Array
 64 |     window_index:int
 65 |     max_steps_in_episode: int
 66 |     start_index: int # This should be here because its the same for all agents, but it changes for all agents when resetting (this is why its not in Params)
 67 |     step_counter: int
 68 |     
 69 | 
 70 | 
 71 | @struct.dataclass
 72 | class WorldState(LoadedEnvState):
 73 |     # But everything here that is not loaded from the base config but shared by all agents
 74 |     best_bids: jnp.ndarray
 75 |     best_asks: jnp.ndarray
 76 |     time: chex.Array
 77 |     order_id_counter: int
 78 |     mid_price:float
 79 |     delta_time: float
 80 | 
 81 | 
 82 | # Define a combined (multi–agent) state that extends the base order book state
 83 | @struct.dataclass
 84 | class MultiAgentState():
 85 |     # Sub–state for market maker and execution agent.
 86 |     world_state: WorldState
 87 | 
 88 |     agent_states: list[Any]
 89 | 
 90 | 
 91 | 
 92 | @struct.dataclass
 93 | class MMEnvState():
 94 |     inventory: int
 95 |     total_PnL: float
 96 |     cash_balance: float
 97 | 
 98 | 
 99 | @struct.dataclass
100 | class ExecEnvState():
101 |     init_price: int
102 |     task_to_execute: int
103 |     quant_executed: int
104 |     # rewards
105 |     total_revenue: float
106 |     drift_return: float
107 |     advantage_return: float
108 |     slippage_rm: float
109 |     price_adv_rm: float
110 |     price_drift_rm: float
111 |     vwap_rm: float
112 |     is_sell_task: int
113 |     trade_duration: float
114 | 
115 | 
116 | 
117 | 
118 | 
119 | 
120 | ########################################################################################
121 | ########################################################################################
122 | # Params
123 | ########################################################################################
124 | ########################################################################################
125 | 
126 | @struct.dataclass
127 | class LoadedEnvParams:
128 |     message_data: chex.Array
129 |     book_data: chex.Array
130 |     init_states_array: chex.Array
131 | 
132 | 
133 | 
134 | # Define a combined parameters class.
135 | # Logic: All the data is in BaseParams. All the things that depend on all agents are added to it (e.g. num_msgs_per_step). The rest stays in the config
136 | @struct.dataclass
137 | class MultiAgentParams():
138 |     loaded_params: LoadedEnvParams
139 | 
140 |     # Put everything here that is shared by all agents, and will be determined by the world config (if its hard encoded but it in world config)
141 |     # However put static things on the self if they have to be calculated (num_msg_per_step) or in the config if not
142 |     #num_msgs_per_step: int
143 | 
144 |     agent_params: list[Any] # List of either MMEnvParams or ExecEnvParams
145 | 
146 | 
147 | @struct.dataclass
148 | class MMEnvParams():
149 |     trader_id: chex.Array
150 |     time_delay_obs_act: chex.Array
151 |     normalize: chex.Array
152 | 
153 | 
154 | 
155 | @struct.dataclass
156 | class ExecEnvParams():
157 |     trader_id: chex.Array
158 |     task_size: chex.Array 
159 |     reward_lambda: chex.Array
160 |     time_delay_obs_act: chex.Array
161 |     normalize: chex.Array
162 | 
163 | 
164 | 


--------------------------------------------------------------------------------
/gymnax_exchange/jaxrl/MARL/baseline_eval/simple_triangular_plot.py:
--------------------------------------------------------------------------------
  1 | #!/usr/bin/env python3
  2 | """
  3 | Simple script to create triangular heatmap showing MarketMaker and ExecutionAgent rewards
  4 | for the 4 combinations: BB, BL, LB, LL
  5 | """
  6 | 
  7 | import pandas as pd
  8 | import numpy as np
  9 | import matplotlib.pyplot as plt
 10 | import matplotlib.patches as mpatches
 11 | import matplotlib.colors as mcolors
 12 | import matplotlib.lines as mlines
 13 | import argparse
 14 | import sys
 15 | 
 16 | def load_data(csv_path):
 17 |     """Load data from CSV file"""
 18 |     try:
 19 |         df = pd.read_csv(csv_path)
 20 |         print("Loaded data:")
 21 |         print(df)
 22 |         return df
 23 |     except FileNotFoundError:
 24 |         print(f"Error: File '{csv_path}' not found.")
 25 |         sys.exit(1)
 26 |     except Exception as e:
 27 |         print(f"Error loading file: {e}")
 28 |         sys.exit(1)
 29 | 
 30 | def create_triangular_plot(df, reward_label='Reward', save_path=None):
 31 |     """Create triangular plot with clear mapping"""
 32 |     
 33 |     # Use mean aggregation by default for multiple episodes
 34 |     agg_data = df.groupby('Configuration')[['MarketMaker', 'ExecutionAgent']].mean()
 35 |     
 36 |     print(f"Aggregated data (using mean):")
 37 |     print(agg_data)
 38 |     
 39 |     # Create a 2x2 grid
 40 |     fig, ax = plt.subplots(1, 1, figsize=(10, 8))
 41 |     
 42 |     # Define the grid positions for each configuration
 43 |     # We want: BB bottom-left, BL bottom-right, LB top-left, LL top-right
 44 |     positions = {
 45 |         'Config_BB': (1, 0),  # bottom-left (row 1, col 0)
 46 |         'Config_BL': (1, 1),  # bottom-right (row 1, col 1) 
 47 |         'Config_LB': (0, 0),  # top-left (row 0, col 0)
 48 |         'Config_LL': (0, 1)   # top-right (row 0, col 1)
 49 |     }
 50 |     
 51 |     # Color settings
 52 |     vmin, vmax = -40, 0
 53 |     cmap = mcolors.LinearSegmentedColormap.from_list('RdYlGn', ['red', 'yellow', 'green'])
 54 |     norm = mcolors.Normalize(vmin=vmin, vmax=vmax)
 55 |     
 56 |     # Process each configuration
 57 |     for config in agg_data.index:
 58 |         mm_reward = agg_data.loc[config, 'MarketMaker']
 59 |         ea_reward = agg_data.loc[config, 'ExecutionAgent']
 60 |         
 61 |         print(f"Processing {config}: MM={mm_reward:.1f}, EA={ea_reward:.1f}")
 62 |         
 63 |         if config in positions:
 64 |             grid_row, grid_col = positions[config]
 65 |             
 66 |             # Clamp values for color mapping
 67 |             mm_color = np.clip(mm_reward, vmin, vmax)
 68 |             ea_color = np.clip(ea_reward, vmin, vmax)
 69 |             
 70 |             # Create triangular patches
 71 |             # Upper triangle for MarketMaker
 72 |             upper_triangle = mpatches.Polygon(
 73 |                 [(grid_col, grid_row+1), (grid_col+1, grid_row+1), (grid_col+1, grid_row)], 
 74 |                 closed=True, 
 75 |                 facecolor=cmap(norm(mm_color)), 
 76 |                 edgecolor='black', 
 77 |                 linewidth=1
 78 |             )
 79 |             ax.add_patch(upper_triangle)
 80 |             
 81 |             # Lower triangle for ExecutionAgent
 82 |             lower_triangle = mpatches.Polygon(
 83 |                 [(grid_col, grid_row), (grid_col, grid_row+1), (grid_col+1, grid_row)],
 84 |                 closed=True, 
 85 |                 facecolor=cmap(norm(ea_color)),
 86 |                 edgecolor='black', 
 87 |                 linewidth=1
 88 |             )
 89 |             ax.add_patch(lower_triangle)
 90 |             
 91 |             # Add labels and values for both triangles
 92 |             # Upper triangle: MM label and value
 93 |             ax.text(grid_col+0.75, grid_row+0.85, 'MM', 
 94 |                    ha='center', va='center', fontsize=18, fontweight='bold', color='black')
 95 |             ax.text(grid_col+0.75, grid_row+0.65, f'{mm_reward:.1f}', 
 96 |                    ha='center', va='center', fontsize=18, fontweight='bold', color='black')
 97 |             
 98 |             # Lower triangle: EXEC label and value  
 99 |             ax.text(grid_col+0.25, grid_row+0.35, 'EXEC', 
100 |                    ha='center', va='center', fontsize=18, fontweight='bold', color='black')
101 |             ax.text(grid_col+0.25, grid_row+0.15, f'{ea_reward:.1f}', 
102 |                    ha='center', va='center', fontsize=18, fontweight='bold', color='black')
103 |             
104 |             print(f"  -> Placed at grid position ({grid_row}, {grid_col})")
105 |     
106 |     # Configure the plot
107 |     ax.set_xlim(0, 2)
108 |     ax.set_ylim(0, 2)
109 |     ax.set_aspect('equal')
110 |     
111 |     # Set ticks and labels
112 |     ax.set_xticks([0.5, 1.5])
113 |     ax.set_xticklabels(['Baseline', 'Learned'], fontsize=18)
114 |     ax.set_xlabel('ExecutionAgent (Slippage)', fontsize=24)
115 |     
116 |     ax.set_yticks([0.5, 1.5])
117 |     ax.set_yticklabels(['Learned', 'Baseline'],fontsize=18,rotation='vertical')  # Note: y-axis is inverted in matplotlib
118 |     ax.set_ylabel('MarketMaker (Portfolio Value)', fontsize=24)
119 |     
120 |     # Invert y-axis to make bottom-left = BB
121 |     ax.invert_yaxis()
122 |     
123 |     # Add title
124 |     ax.set_title(f'{reward_label}: Baseline vs Learned', 
125 |                 fontsize=24, pad=20)
126 |     
127 |     # Add colorbar
128 |     sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
129 |     sm.set_array([])
130 |     cbar = plt.colorbar(sm, ax=ax, shrink=0.8)
131 |     cbar.set_label(f'{reward_label}', rotation=270, labelpad=15, fontsize=18)
132 |     
133 |     plt.tight_layout()
134 |     
135 |     if save_path:
136 |         plt.savefig(save_path, dpi=300, bbox_inches='tight')
137 |         print(f"Plot saved to: {save_path}")
138 |     
139 |     return fig
140 | 
141 | def main():
142 |     # Set up command line argument parsing
143 |     parser = argparse.ArgumentParser(description='Create triangular heatmap for agent rewards')
144 |     parser.add_argument('csv_file', help='Path to CSV file containing the data')
145 |     parser.add_argument('--reward-label', default='Rewards', 
146 |                        help='Label for the reward type being displayed (default: Rewards)')
147 |     parser.add_argument('--save-path', help='Path to save the plot (optional)')
148 |     
149 |     args = parser.parse_args()
150 |     
151 |     # Load the data
152 |     df = load_data(args.csv_file)
153 |     
154 |     # Create the plot
155 |     save_path = args.save_path
156 |     if save_path is None:
157 |         # Generate default save path
158 |         import os
159 |         base_name = os.path.splitext(os.path.basename(args.csv_file))[0]
160 |         save_path = f'{base_name}_triangular.png'
161 |     
162 |     fig = create_triangular_plot(df, reward_label=args.reward_label, save_path=save_path)
163 |     
164 |     print("Plot created and saved successfully!")
165 | 
166 | if __name__ == "__main__":
167 |     main()
168 | 


--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
  1 | absl-py==2.1.0
  2 | aiohappyeyeballs==2.6.1
  3 | aiohttp==3.12.14
  4 | aiosignal==1.4.0
  5 | alabaster==1.0.0
  6 | annotated-types==0.7.0
  7 | antlr4-python3-runtime==4.9.3
  8 | anyio==4.8.0
  9 | argon2-cffi==23.1.0
 10 | argon2-cffi-bindings==21.2.0
 11 | arrow==1.3.0
 12 | astroid==3.3.8
 13 | asttokens==3.0.0
 14 | async-lru==2.0.4
 15 | attrs==25.1.0
 16 | autopep8==2.3.2
 17 | babel==2.17.0
 18 | beautifulsoup4==4.13.3
 19 | black==25.1.0
 20 | bleach==6.2.0
 21 | blinker==1.9.0
 22 | brax==0.12.1
 23 | build==1.2.2.post1
 24 | cachetools==5.5.2
 25 | certifi==2025.1.31
 26 | cffi==1.17.1
 27 | charset-normalizer==3.4.1
 28 | cheroot==10.0.1
 29 | chex==0.1.88
 30 | click==8.1.8
 31 | cloudpickle==3.1.1
 32 | codespell==2.4.1
 33 | colorama==0.4.6
 34 | comm==0.2.2
 35 | contourpy==1.3.1
 36 | cycler==0.12.1
 37 | dataclasses-json==0.6.7
 38 | debugpy==1.8.12
 39 | decorator==5.2.1
 40 | defusedxml==0.7.1
 41 | dill==0.3.9
 42 | distrax==0.1.5
 43 | dm-env==1.6
 44 | dm-tree==0.1.9
 45 | docker-pycreds==0.4.0
 46 | docstring_parser==0.16
 47 | docutils==0.21.2
 48 | einops==0.8.1
 49 | etils==1.12.0
 50 | exceptiongroup==1.2.2
 51 | executing==2.2.0
 52 | Farama-Notifications==0.0.4
 53 | fastjsonschema==2.21.1
 54 | filelock==3.17.0
 55 | flake8==7.1.2
 56 | Flask==3.1.0
 57 | flask-cors==5.0.1
 58 | # Editable Git install with no remote (flax==0.10.2)
 59 | -e /opt/flax
 60 | fonttools==4.56.0
 61 | fqdn==1.5.1
 62 | frozenlist==1.7.0
 63 | fsspec==2025.2.0
 64 | furo==2024.8.6
 65 | gast==0.6.0
 66 | gcsfs==2025.2.0
 67 | gdown==5.2.0
 68 | gitdb==4.0.12
 69 | GitPython==3.1.44
 70 | glfw==2.8.0
 71 | google-api-core==2.25.1
 72 | google-auth==2.40.3
 73 | google-auth-oauthlib==1.2.2
 74 | google-cloud-core==2.4.3
 75 | google-cloud-storage==3.2.0
 76 | google-crc32c==1.7.1
 77 | google-resumable-media==2.7.2
 78 | googleapis-common-protos==1.70.0
 79 | greenlet==3.1.1
 80 | grpcio==1.70.0
 81 | gviz-api==1.10.0
 82 | gym==0.26.2
 83 | gym-notices==0.0.8
 84 | gymnasium==1.0.0
 85 | gymnax==0.0.8
 86 | h11==0.14.0
 87 | httpcore==1.0.7
 88 | httpx==0.28.1
 89 | huggingface-hub==0.29.1
 90 | humanize==4.12.1
 91 | hydra-core==1.3.2
 92 | idna==3.10
 93 | imagesize==1.4.1
 94 | importlib_metadata==8.6.1
 95 | importlib_resources==6.5.2
 96 | iniconfig==2.0.0
 97 | ipdb==0.13.13
 98 | ipykernel==6.29.5
 99 | ipython==8.32.0
100 | ipywidgets==8.1.5
101 | isoduration==20.11.0
102 | isort==6.0.0
103 | itsdangerous==2.2.0
104 | jaraco.functools==4.2.1
105 | jax==0.6.1
106 | jax-cuda12-pjrt==0.6.1
107 | jax-cuda12-plugin==0.6.1
108 | jaxlib==0.6.1
109 | jaxopt==0.8.5
110 | jaxtyping==0.2.36
111 | jedi==0.19.2
112 | Jinja2==3.1.5
113 | json5==0.10.0
114 | jsonpointer==3.0.0
115 | jsonschema==4.23.0
116 | jsonschema-specifications==2024.10.1
117 | jupyter==1.1.1
118 | jupyter-cache==1.0.1
119 | jupyter-console==6.6.3
120 | jupyter-events==0.12.0
121 | jupyter-lsp==2.2.5
122 | jupyter_client==8.6.3
123 | jupyter_core==5.7.2
124 | jupyter_server==2.15.0
125 | jupyter_server_terminals==0.5.3
126 | jupyterlab==4.3.5
127 | jupyterlab_pygments==0.3.0
128 | jupyterlab_server==2.27.3
129 | jupyterlab_widgets==3.0.13
130 | kiwisolver==1.4.8
131 | Markdown==3.8.2
132 | markdown-it-py==3.0.0
133 | MarkupSafe==3.0.2
134 | marshmallow==3.25.1
135 | matplotlib==3.10.0
136 | matplotlib-inline==0.1.7
137 | mccabe==0.7.0
138 | mdit-py-plugins==0.4.2
139 | mdurl==0.1.2
140 | mistune==3.1.2
141 | ml_collections==1.0.0
142 | ml_dtypes==0.5.1
143 | more-itertools==10.7.0
144 | mpmath==1.3.0
145 | msgpack==1.1.0
146 | mujoco==3.2.7
147 | mujoco-mjx==3.2.7
148 | multidict==6.6.3
149 | mypy-extensions==1.0.0
150 | myst-nb==1.2.0
151 | myst-parser==4.0.1
152 | Navix==0.7.0
153 | nbclient==0.10.2
154 | nbconvert==7.16.6
155 | nbformat==5.10.4
156 | nest-asyncio==1.6.0
157 | networkx==3.4.2
158 | notebook==7.3.2
159 | notebook_shim==0.2.4
160 | # Editable install with no version control (nsys-jax==0.1.dev1085+gb0ec72a)
161 | -e /opt/nsys-jax/.github/container/nsys_jax
162 | numpy==1.26.4
163 | nvidia-cublas-cu12==12.6.4.1
164 | nvidia-cuda-cupti-cu12==12.6.80
165 | nvidia-cuda-nvcc-cu12==12.9.86
166 | nvidia-cuda-nvrtc-cu12==12.6.77
167 | nvidia-cuda-runtime-cu12==12.6.77
168 | nvidia-cudnn-cu12==9.11.0.98
169 | nvidia-cufft-cu12==11.3.0.4
170 | nvidia-cufile-cu12==1.11.1.6
171 | nvidia-curand-cu12==10.3.7.77
172 | nvidia-cusolver-cu12==11.7.1.2
173 | nvidia-cusparse-cu12==12.5.4.2
174 | nvidia-cusparselt-cu12==0.6.3
175 | nvidia-nccl-cu12==2.26.2
176 | nvidia-nvjitlink-cu12==12.6.85
177 | nvidia-nvshmem-cu12==3.3.9
178 | nvidia-nvtx-cu12==12.6.77
179 | oauthlib==3.3.1
180 | omegaconf==2.3.0
181 | opt_einsum==3.4.0
182 | optax==0.2.4
183 | orbax==0.1.9
184 | orbax-checkpoint==0.11.6
185 | orbax-export==0.0.6
186 | overrides==7.7.0
187 | packaging==24.2
188 | pandas==2.2.3
189 | pandocfilters==1.5.1
190 | parso==0.8.4
191 | pathspec==0.12.1
192 | pexpect==4.9.0
193 | pillow==11.1.0
194 | pip-tools==7.4.1
195 | platformdirs==4.3.6
196 | pluggy==1.5.0
197 | pockets==0.9.1
198 | prometheus_client==0.21.1
199 | prompt_toolkit==3.0.50
200 | propcache==0.3.2
201 | proto-plus==1.26.1
202 | protobuf==5.29.3
203 | psutil==7.0.0
204 | ptyprocess==0.7.0
205 | pure_eval==0.2.3
206 | pyarrow==19.0.0
207 | pyasn1==0.6.1
208 | pyasn1_modules==0.4.2
209 | pycodestyle==2.12.1
210 | pycparser==2.22
211 | pydantic==2.10.6
212 | pydantic_core==2.27.2
213 | pyflakes==3.2.0
214 | Pygments==2.19.1
215 | pylint==3.3.4
216 | PyOpenGL==3.1.9
217 | pyparsing==3.2.1
218 | pyproject_hooks==1.2.0
219 | PySocks==1.7.1
220 | pytest==8.3.4
221 | python-dateutil==2.9.0.post0
222 | python-json-logger==3.2.1
223 | pytinyrenderer==0.0.14
224 | pytz==2025.1
225 | PyYAML==6.0.2
226 | pyzmq==26.2.1
227 | referencing==0.36.2
228 | regex==2024.11.6
229 | requests==2.32.3
230 | requests-oauthlib==2.0.0
231 | rfc3339-validator==0.1.4
232 | rfc3986-validator==0.1.1
233 | rich==13.9.4
234 | rlax==0.1.6
235 | rpds-py==0.23.1
236 | rsa==4.9.1
237 | rwkv==0.8.28
238 | safetensors==0.5.3
239 | scipy==1.15.2
240 | seaborn==0.13.2
241 | Send2Trash==1.8.3
242 | sentry-sdk==2.22.0
243 | setproctitle==1.3.5
244 | setuptools==75.8.0
245 | setuptools-scm==8.2.0
246 | shtab==1.7.1
247 | simplejson==3.20.1
248 | six==1.17.0
249 | smmap==5.0.2
250 | sniffio==1.3.1
251 | snowballstemmer==2.2.0
252 | soupsieve==2.6
253 | Sphinx==8.1.3
254 | sphinx-autobuild==2024.10.3
255 | sphinx-autodoc-typehints==3.0.1
256 | sphinx-basic-ng==1.0.0b2
257 | sphinx-copybutton==0.5.2
258 | sphinx-github-changelog==1.4.0
259 | sphinx-rtd-theme==3.0.2
260 | sphinxcontrib-applehelp==2.0.0
261 | sphinxcontrib-devhelp==2.0.0
262 | sphinxcontrib-htmlhelp==2.1.0
263 | sphinxcontrib-jquery==4.1
264 | sphinxcontrib-jsmath==1.0.1
265 | sphinxcontrib-napoleon==0.7
266 | sphinxcontrib-qthelp==2.0.0
267 | sphinxcontrib-serializinghtml==2.0.0
268 | SQLAlchemy==2.0.38
269 | stack-data==0.6.3
270 | starlette==0.46.0
271 | sympy==1.14.0
272 | tabulate==0.9.0
273 | tensorboard==2.20.0
274 | tensorboard-data-server==0.7.2
275 | tensorboard_plugin_profile==2.19.9
276 | tensorboardX==2.6.4
277 | tensorflow-probability==0.25.0
278 | tensorstore==0.1.71
279 | terminado==0.18.1
280 | tinycss2==1.4.0
281 | tokenizers==0.21.0
282 | tomli==2.2.1
283 | tomlkit==0.13.2
284 | toolz==1.0.0
285 | #torch==2.7.1
286 | tornado==6.4.2
287 | tqdm==4.67.1
288 | traitlets==5.14.3
289 | transformer_engine @ file:///opt/transformer-engine/dist/transformer_engine-1.14.0+87fbe812f-cp312-cp312-linux_x86_64.whl#sha256=9bb7e91c6687050a3dc5d3b80d5719df293da899c238057beda3ee0f55e3253a
290 | transformers==4.49.0
291 | treescope==0.1.9
292 | trimesh==4.6.3
293 | triton==3.3.1
294 | typeguard==4.4.2
295 | types-python-dateutil==2.9.0.20241206
296 | typing-inspect==0.9.0
297 | typing_extensions==4.12.2
298 | tyro==0.9.16
299 | tzdata==2025.1
300 | uncertainties==3.2.2
301 | uri-template==1.3.0
302 | urllib3==2.3.0
303 | uvicorn==0.34.0
304 | wandb==0.19.7
305 | watchfiles==1.0.4
306 | wcwidth==0.2.13
307 | webcolors==24.11.1
308 | webencodings==0.5.1
309 | websocket-client==1.8.0
310 | websockets==15.0
311 | Werkzeug==3.1.3
312 | wheel==0.45.1
313 | widgetsnbextension==4.0.13
314 | wrapt==1.17.2
315 | xprof==2.19.9
316 | yarl==1.20.1
317 | zipp==3.21.0
318 | 


--------------------------------------------------------------------------------
/gymnax_exchange/jaxob/jaxob_config.py:
--------------------------------------------------------------------------------
  1 | import gymnax_exchange.jaxob.jaxob_constants as cst
  2 | import os
  3 | from typing import OrderedDict, Tuple,  Literal,Union,List
  4 | 
  5 | from dataclasses import dataclass,field
  6 | 
  7 | 
  8 | 
  9 | 
 10 | 
 11 | @dataclass(frozen=True)
 12 | class JAXLOB_Configuration:
 13 |     maxint : int = cst.MaxInt._64_Bit_Signed.value
 14 |     init_id :int = cst.INITID
 15 |     cancel_mode: int= cst.CancelMode.INCLUDE_INITS.value
 16 |     seed: int =cst.SEED
 17 |     nTrades : int=cst.NTRADE_CAP
 18 |     nOrders : int =cst.NORDER_CAP
 19 |     simulator_mode=cst.SimulatorMode.GENERAL_EXCHANGE.value
 20 |     empty_slot_val=cst.EMPTY_SLOT
 21 |     debug_mode:bool=False
 22 |     start_resolution: int = 50  # Episodes from data start every n seconds.
 23 |     alphatradePath: str = os.path.expanduser("~")
 24 |     dataPath: str = os.path.expanduser("~")+"/data"
 25 |     stock: str = "AMZN"
 26 |     timePeriod: str = "2017Jan_oneday" # Needs to be the appropriate directory name. "2017Jan_oneday", "2024"
 27 | 
 28 | 
 29 | @dataclass(frozen=True)
 30 | class MarketMaking_EnvironmentConfig():
 31 |     # action_space options: "fixed_prices", "fixed_quants", "AvSt", "spread_skew", "directional_trading", "simple"
 32 |     action_space: str = "spread_skew"
 33 | 
 34 |     # observation_space options: "engineered", "messages", "messages_new_tokenizer", "basic"
 35 |     observation_space: str = "engineered"
 36 |     #end_fn: Literal["force_market_order", "unwind_ref_price","do_nothing"] = "unwind_ref_price"
 37 |     # Values for spread skew action space
 38 |     spread_multiplier: float = 3.0 #50.0
 39 |     skew_multiplier: float = 5.0 #100.0
 40 |     n_ticks_in_book : int = 1
 41 |     num_messages_by_agent:int=4
 42 |     num_action_messages_by_agent=2 # will be set automcatically down below
 43 |     fixed_quant_value:int=10
 44 |     n_actions: int = 8 # Only used for fixed_prices
 45 |     debug_mode:bool=False
 46 |     time_delay_obs_act:int=0
 47 |     normalize:bool=True
 48 |     short_name:str="MM" # For agent naming e.g. in the obs dict
 49 |     seconds_before_episode_end:int=5
 50 |     # Fixed action settings
 51 |     fixed_action_setting: bool = False
 52 |     fixed_action: int = 0
 53 |     #Control for simple market making action space
 54 |     sell_buy_all_option: bool= False #Whether selling the entire inventory is possible
 55 |     simple_nothing_action: bool = True # Whether or not the simple action space has a nothing action
 56 | 
 57 |     # Reward
 58 |     inv_penalty: str = "none"  # options: "none", "linear", "quadratic", "threshold"
 59 |     reward_space: str = "buy_sell_pnl"  # options: "zero_inv", "pnl", "buy_sell_pnl", "complex", "portfolio_value", "portfolio_value_scaled", "spooner", "spooner_damped", "spooner_scaled", "delta_netWorth","weight_pnl_inventory_pnl"
 60 |     reference_price_portfolio_value: str = "mid"  # options: "mid", "best_bid_ask", "near_touch"
 61 |     inv_penalty_lambda: float = 1.0
 62 |     inv_penalty_quadratic_factor: float = 50.0 #Represents N for penalty = 1/N * (inv ** 2) if quadratic penalty is used
 63 |     multiplier_type: str = "tick" # options:  "tick" #DO NOT USE "spread" it is WRONG. 
 64 |     clip_reward: bool = False
 65 |     based_on_mid_price_of_action: bool = True
 66 |     exclude_extreme_spreads: bool= False
 67 |     # Weights for complex reward function:
 68 |     inventoryPnL_lambda: float = 0.5
 69 |     unrealizedPnL_lambda: float = 0.1
 70 |     asymmetrically_dampened_lambda: float = 0.8
 71 | 
 72 | 
 73 |     def __post_init__(self):
 74 |         # Since the class is frozen, we need to use object.__setattr__ to modify n_actions
 75 |         # Number of messages includes action messages and cancel messages!
 76 |         if self.action_space == "fixed_quants":
 77 |             object.__setattr__(self, 'n_actions', 8)
 78 |             object.__setattr__(self, 'num_messages_by_agent', 4)
 79 |             object.__setattr__(self, 'num_action_messages_by_agent', 2)
 80 |         elif self.action_space == "spread_skew":
 81 |             object.__setattr__(self, 'n_actions', 6)
 82 |             object.__setattr__(self, 'num_messages_by_agent', 4)
 83 |             object.__setattr__(self, 'num_action_messages_by_agent', 2)
 84 |         elif self.action_space == "directional_trading":
 85 |             object.__setattr__(self, 'n_actions', 3)
 86 |             object.__setattr__(self, 'num_messages_by_agent', 4)
 87 |             object.__setattr__(self, 'num_action_messages_by_agent', 2)
 88 |         elif self.action_space == "AvSt":
 89 |             object.__setattr__(self, 'n_actions', 8)
 90 |             object.__setattr__(self, 'num_messages_by_agent', 4)
 91 |             object.__setattr__(self, 'num_action_messages_by_agent', 2)
 92 |         elif self.action_space == "fixed_prices":
 93 |             object.__setattr__(self, 'num_messages_by_agent', self.n_actions*2)
 94 |             object.__setattr__(self, 'num_action_messages_by_agent', self.n_actions)
 95 |         elif self.action_space == "simple":
 96 |             object.__setattr__(self, 'n_actions', 4)
 97 |             object.__setattr__(self, 'num_messages_by_agent', 4)
 98 |             object.__setattr__(self, 'num_action_messages_by_agent', 2)
 99 | 
100 | 
101 | @dataclass(frozen=True)
102 | class Execution_EnvironmentConfig():
103 |     n_ticks_in_book : int = 1
104 |     task: str = "random"  # options: "random", "buy", "sell"
105 |     action_type: str = "pure"  # options: "delta", "pure"
106 |     action_space: str = "fixed_quants_complex"  # options: "fixed_quants", "fixed_prices", "fixed_quants_complex", "simplest_case", "fixed_quants_1msg"
107 |     observation_space: str = "engineered"  # options: "engineered", "basic", "simplest_case"
108 |     reward_space: str = "normal"  # options: "normal", "finish_fast", "simplest_case"
109 |     #end_fn:Literal["force_market_order","unwind_FT"]="unwind_FT"
110 |     task_size:int= 500
111 |     n_actions:int=5 # will be set automatically in the post init function
112 |     fixed_quant_value:int=10
113 |     num_messages_by_agent:int=8 # will be set automatically in the post init function
114 |     num_action_messages_by_agent:int=4 # will be set automatically in the post init function
115 |     reward_lambda:float= 0.0
116 |     time_delay_obs_act:int=0
117 |     debug_mode:bool=False
118 |     normalize:bool=True
119 |     short_name:str="EXE"
120 |     seconds_before_episode_end:int=5
121 |     doom_price_penalty: float = 0.1
122 |     larger_far_touch_quant: bool = False
123 |     
124 | 
125 | 
126 |     def __post_init__(self):
127 |         # Since the class is frozen, we need to use object.__setattr__ to modify n_actions
128 |         # Number of messages includes action messages and cancel messages!
129 |         if self.action_space == "fixed_quants":
130 |             object.__setattr__(self, 'n_actions', 5)
131 |             object.__setattr__(self, 'num_messages_by_agent', 8)
132 |             object.__setattr__(self, 'num_action_messages_by_agent', 4)
133 |         elif self.action_space == "fixed_prices":
134 |             object.__setattr__(self, 'num_messages_by_agent', self.n_actions*2)
135 |             object.__setattr__(self, 'num_action_messages_by_agent', self.n_actions)
136 |         elif self.action_space == "fixed_quants_complex":
137 |             object.__setattr__(self, 'n_actions', 13)
138 |             object.__setattr__(self, 'num_messages_by_agent', 8)
139 |             object.__setattr__(self, 'num_action_messages_by_agent', 4)
140 |         elif self.action_space == "simplest_case":
141 |             object.__setattr__(self, 'n_actions', 3)
142 |             object.__setattr__(self, 'num_messages_by_agent', 4) # Includes cancel messages
143 |             object.__setattr__(self, 'num_action_messages_by_agent', 2)
144 |         elif self.action_space == "fixed_quants_1msg":
145 |             object.__setattr__(self, 'n_actions', 5)
146 |             object.__setattr__(self, 'num_messages_by_agent', 2)
147 |             object.__setattr__(self, 'num_action_messages_by_agent', 1)
148 |         elif self.action_space == "twap":
149 |             object.__setattr__(self, 'n_actions', 1)
150 |             object.__setattr__(self, 'num_messages_by_agent', 4)
151 |             object.__setattr__(self, 'num_action_messages_by_agent', 2)
152 | 
153 | 
154 | 
155 | 
156 | @dataclass(frozen=True)
157 | class World_EnvironmentConfig(JAXLOB_Configuration):
158 |     n_data_msg_per_step: int = 1
159 |     window_selector = -1 # -1 means random window
160 |     ep_type :str = "fixed_steps" # fixed_steps, fixed_time
161 |     episode_time: int = 50 # counted by seconds, 1800s=0.5h or steps
162 |     day_start = 34200  # 09:30
163 |     day_end = 57600  # 16:00
164 |     nOrdersPerSide=100 #100
165 |     nTradesLogged=100
166 |     book_depth=10
167 |     n_ticks_in_book = 10 # Depth of PP actions
168 |     customIDCounter=0
169 |     tick_size=100
170 |     trader_id_range_start=-100 # -1 is reserved for the placeholder in the messages object
171 |     placeholder_order_id = -9
172 |     last_step_seconds = 5
173 |     artificial_trader_id_end_episode = -666666 # Artificial trader id for the trade that is artifically added at the end of the episode (this is not really used)
174 |     artificial_order_id_end_episode = -666666 # Artificial order id for the trade that is artifically added at the end of the episode (this is not really used)
175 |     debug_mode:bool=False
176 |     any_message_obs_space:bool=False # TODO: set this automatically in a post init function based on the obs spaces of each agent type
177 |     order_id_counter_start_when_resetting:int=-200
178 |     shuffle_action_messages:bool=True
179 |     use_pickles_for_init:bool= True
180 |     save_raw_observations:bool=False
181 | 
182 | 
183 | @dataclass(frozen=True)
184 | class MultiAgentConfig():
185 |     #world_config: World_EnvironmentConfig = field(default_factory=lambda: World_EnvironmentConfig())
186 |     world_config: World_EnvironmentConfig = World_EnvironmentConfig()
187 | 
188 |     # list_of_agents_configs: List = field(default_factory=lambda: [
189 |     #     MarketMaking_EnvironmentConfig(),
190 |     #     Execution_EnvironmentConfig()
191 |     # ])
192 | 
193 |     dict_of_agents_configs: dict = field(default_factory=lambda: dict([
194 |         ("MarketMaking", MarketMaking_EnvironmentConfig()),
195 |         ("Execution", Execution_EnvironmentConfig())
196 |     ]))
197 |     number_of_agents_per_type: list = field(default_factory=lambda: [1,1]) # This is only the default value, we change it in the yaml RL file
198 | 
199 | 
200 |     # list_of_agents_configs = [
201 |     # list_of_agents_configs: List =field(default_factory=lambda :[Execution_EnvironmentConfig()])
202 |     #     MarketMaking_EnvironmentConfig(),
203 |     #     #Execution_EnvironmentConfig(),
204 |     #     #MarketMaking_EnvironmentConfig(),
205 |     # ]
206 |     # number_of_agents_per_type = [2]
207 | 
208 | 
209 | 


--------------------------------------------------------------------------------
/gymnax_exchange/jaxob/jorderbook.py:
--------------------------------------------------------------------------------
  1 | from functools import partial
  2 | import importlib
  3 | from os import remove
  4 | from readline import remove_history_item
  5 | from typing import Dict, NamedTuple, Optional, Tuple
  6 | from unicodedata import bidirectional
  7 | import gymnax_exchange.jaxob.JaxOrderBookArrays as job
  8 | from gymnax_exchange.jaxob.jaxob_config import Configuration
  9 | job = importlib.reload(job)
 10 | import jax
 11 | from jax import numpy as jnp
 12 | from jax import lax
 13 | import chex
 14 | 
 15 | 
 16 | 
 17 | class LobState(NamedTuple):
 18 |     asks: jnp.ndarray
 19 |     bids: jnp.ndarray
 20 |     trades: jnp.ndarray
 21 |     key : chex.PRNGKey
 22 | 
 23 | 
 24 | 
 25 | class OrderBook():
 26 |     def __init__(
 27 |             self: 'OrderBook',
 28 |             cfg: Optional[Configuration] = None,
 29 |         ) -> None:
 30 |         self.cfg = cfg if cfg is not None else Configuration()
 31 | 
 32 |     @jax.jit
 33 |     def init(self: 'OrderBook') -> LobState:
 34 |         asks = (jnp.ones((self.cfg.nOrders, 6)) * -1).astype(jnp.int32)
 35 |         bids = (jnp.ones((self.cfg.nOrders, 6)) * -1).astype(jnp.int32)
 36 |         trades = (jnp.ones((self.cfg.nTrades, 8)) * -1).astype(jnp.int32)
 37 |         key=jax.random.PRNGKey(self.cfg.seed)
 38 |         return LobState(asks, bids, trades,key)
 39 | 
 40 |     @jax.jit
 41 |     def reset(
 42 |             self: 'OrderBook',
 43 |             l2_book: Optional[jnp.ndarray] = None,
 44 |             time: Optional[jnp.ndarray] = None,
 45 |         ) -> LobState:
 46 |         """"""
 47 |         state = self.init()
 48 |         if l2_book is not None:
 49 |             if time is None:
 50 |                 time=jnp.array([0,0])
 51 |             msgs = job.init_msgs_from_l2(self.cfg,l2_book,time=time)
 52 |             state = self.process_orders_array(state, msgs)
 53 |         return state
 54 | 
 55 |     def process_order(
 56 |             self: 'OrderBook',
 57 |             state: LobState,
 58 |             quote: Dict,
 59 |             from_data: bool = False,
 60 |             verbose: bool = False
 61 |         ) -> LobState:
 62 |         '''Wrapper function for the object class that takes a Dict Object as the quote,
 63 |          ensures the order is conserved and turns the values into a jnp array which is passed to the JNP ProcessOrder function'''
 64 |         #Type, Side,quant,price
 65 |         inttype = 5
 66 |         intside = -1
 67 |         if quote['side'] == 'bid':
 68 |             intside = 1 
 69 | 
 70 |         if quote['type'] == 'limit':
 71 |             inttype = 1
 72 |         elif quote['type'] == 'cancel':
 73 |             inttype = 2
 74 |         elif quote['type'] == 'delete':
 75 |             inttype = 2
 76 |         elif quote['type'] == 'market':
 77 |             inttype = 1
 78 |             intside = intside * -1
 79 | 
 80 |         #TODO: Config so that a type 4 market order is seen as a near touch order. 
 81 | 
 82 |         msg = jnp.array([
 83 |             inttype,
 84 |             intside,
 85 |             quote['quantity'],
 86 |             quote['price'],
 87 |             quote['trade_id'],
 88 |             quote['order_id'],
 89 |             int(quote['timestamp'].split('.')[0]),
 90 |             int(quote['timestamp'].split('.')[1])
 91 |         ], dtype=jnp.int32)
 92 |         (*ordersides,key)=state
 93 |         key,split_key =jax.random.split(key)
 94 |         ordersides, _ = job.cond_type_side(self.cfg,ordersides,(split_key,msg))
 95 |         return LobState(*ordersides,key)
 96 | 
 97 |     @jax.jit
 98 |     def process_order_array(
 99 |             self: 'OrderBook',
100 |             state: LobState,
101 |             quote: jax.Array,
102 |             from_data: bool = False,
103 |             verbose: bool = False,
104 |         ) -> LobState:
105 |         '''Same as process_order but quote is an array.'''
106 |         (*ordersides,key)=state
107 |         key,split_key =jax.random.split(key)
108 |         ordersides, _ = job.cond_type_side(self.cfg,ordersides, (split_key,quote))
109 |         return LobState(*ordersides,key)
110 | 
111 |     @jax.jit
112 |     def process_orders_array(
113 |         self: 'OrderBook',
114 |         state: LobState,
115 |         msgs: jax.Array,
116 |     ) -> LobState:
117 |         '''Wrapper function for the object class that takes a JNP Array of messages (Shape=Nx8), and applies them, in sequence, to the orderbook'''
118 |         (*ordersides,key)=state
119 |         key,split_key =jax.random.split(key)
120 |         return LobState(*job.scan_through_entire_array(self.cfg,split_key,msgs, tuple(ordersides)),key)
121 | 
122 |     @partial(jax.jit, static_argnums=(3,))
123 |     def process_orders_array_l2(
124 |             self: 'OrderBook',
125 |             state: LobState,
126 |             msgs: jax.Array,
127 |             n_levels: int
128 |         ) -> Tuple[LobState, jax.Array]:
129 |         (*ordersides,key)=state
130 |         key,split_key =jax.random.split(key)
131 |         all_asks, all_bids, trades = job.scan_through_entire_array_save_states(self.cfg,split_key,msgs, tuple(ordersides), msgs.shape[0])
132 |         state = LobState(all_asks[-1], all_bids[-1], trades,key)
133 |         # calculate l2 states
134 |         l2_states = jax.vmap(job.get_L2_state, (0, 0, None,None), 0)(all_asks, all_bids, n_levels,self.cfg)
135 |         return state, l2_states
136 | 
137 |     @partial(jax.jit, static_argnums=(2,4))
138 |     def get_volume_at_price(
139 |             self: 'OrderBook',
140 |             state: LobState,
141 |             side: int,
142 |             price: int,
143 |             init_only: bool = False
144 |         ) -> int:
145 | 
146 |         if side == 0:
147 |             side_array = state.asks
148 |         elif side == 1:
149 |             side_array = state.bids
150 |         else:
151 |             raise ValueError('Side must be 0 or 1')
152 | 
153 |         if init_only:
154 |             return job.get_init_volume_at_price(side_array, price,self.cfg)
155 |         else:
156 |             return job.get_volume_at_price(side_array, price)
157 | 
158 |     @jax.jit
159 |     def get_best_price(
160 |             self: 'OrderBook',
161 |             state: LobState,
162 |             side: int
163 |         ) -> int:
164 |         return jax.lax.cond(
165 |             side == 1,
166 |             lambda s: self.get_best_bid(s),
167 |             lambda s: self.get_best_ask(s),
168 |             state
169 |         )
170 |     
171 |     @jax.jit
172 |     def get_best_bid(
173 |             self: 'OrderBook',
174 |             state: LobState
175 |         ) -> int:
176 |         return job.get_best_bid(self.cfg,state.bids)
177 | 
178 |     @jax.jit
179 |     def get_best_ask(
180 |             self: 'OrderBook',
181 |             state: LobState
182 |         ) -> int:
183 |         return job.get_best_ask(self.cfg,state.asks)
184 | 
185 |     @jax.jit
186 |     def get_best_bid_and_ask_inclQuants(
187 |             self: 'OrderBook',
188 |             state: LobState
189 |         ) -> Tuple[jax.Array, jax.Array]:
190 |         return job.get_best_bid_and_ask_inclQuants(self.cfg,state.asks, state.bids)
191 | 
192 |     @partial(jax.jit, static_argnums=(2,))
193 |     def get_L2_state(
194 |             self: 'OrderBook',
195 |             state: LobState,
196 |             n_levels: int
197 |         ) -> jax.Array:
198 |         return job.get_L2_state(state.asks, state.bids, n_levels,self.cfg)
199 |     
200 |     @jax.jit
201 |     def get_side_ids(
202 |             self: 'OrderBook',
203 |             state: LobState,
204 |             side: int
205 |         ) -> jax.Array:
206 |         side_array = jax.lax.cond(
207 |             side == 1,
208 |             lambda a, b: b,
209 |             lambda a, b: a,
210 |             state.asks, state.bids
211 |         )
212 |         return job.get_order_ids(side_array)
213 | 
214 |     @jax.jit
215 |     def get_order(
216 |             self: 'OrderBook',
217 |             state: LobState,
218 |             side: int,
219 |             order_id: int,
220 |             price: Optional[int] = None,
221 |         ) -> jax.Array:
222 |         ''' '''
223 |         side_array = jax.lax.cond(
224 |             side == 1,
225 |             lambda a, b: b,
226 |             lambda a, b: a,
227 |             state.asks, state.bids
228 |         )
229 |         if price is not None:
230 |             return job.get_order_by_id_and_price(side_array, order_id, price)
231 |         else:
232 |             return job.get_order_by_id(side_array, order_id)
233 |         
234 |     @jax.jit
235 |     def get_order_at_time(
236 |             self: 'OrderBook',
237 |             state: LobState,
238 |             side: int,
239 |             time_s: int,
240 |             time_ns: int,
241 |         ) -> jax.Array:
242 |         ''' '''
243 |         side_array = jax.lax.cond(
244 |             side == 1,
245 |             lambda a, b: b,
246 |             lambda a, b: a,
247 |             state.asks, state.bids
248 |         )
249 |         return job.get_order_by_time(side_array, time_s, time_ns)
250 | 
251 |         
252 |     @jax.jit
253 |     def get_next_executable_order(
254 |             self: 'OrderBook',
255 |             state: LobState,
256 |             side: int
257 |         ):
258 |         side_array = jax.lax.cond(
259 |             side == 1,
260 |             lambda a, b: b,
261 |             lambda a, b: a,
262 |             state.asks, state.bids
263 |         )
264 |         return job.get_next_executable_order(self.cfg,side, side_array)
265 |     
266 |     #Flatten and Unflatten functions so that methods can be appropriately jitted. 
267 |     def _tree_flatten(self: 'OrderBook'):
268 |         children = ()  # arrays / dynamic values
269 |         aux_data = {'cfg': self.cfg,}  # static values
270 |         return (children, aux_data)
271 | 
272 |     @classmethod
273 |     def _tree_unflatten(cls, aux_data, children):
274 |         return cls(*children, **aux_data)
275 | 
276 | 
277 | jax.tree_util.register_pytree_node(
278 |     OrderBook,
279 |     OrderBook._tree_flatten,
280 |     OrderBook._tree_unflatten
281 | )
282 | 
283 | 
284 | if __name__ == "__main__":
285 |     print("""Testing functionality of the jax OrderBook object in jorderbook.py 
286 |                 acts as an object wrapper around the JaxOrderBookArrays functions""")
287 |     
288 |     import typing
289 | 
290 |     ob=OrderBook()
291 |     
292 |     dict_quote={'type':'limit','side':'bid','quantity':99,'price':346000,'trade_id':8888,'order_id':8888,'timestamp':'3400.005000000'}
293 | 
294 |     l2init=jnp.array([354200,452,350100,89,361200,100,344000,400,362900,100,343100,100,364000,400,338700,100,371900,1100,337100,1000,372200,100,336400,1000,372300,200,336000,300,372800,1000,333600,1000,374600,1000,332500,100,376700,100,331600,100])
295 |     state=ob.reset(l2init)
296 |     print(state)
297 |     print("Message dict quote processing",ob.process_order(state,dict_quote))
298 |     msg_array=jnp.array([1,1,99,346000,8888,8888,3400,5000000])
299 |     msgs_array=jnp.array([[1,1,99,346000,8888,8888,3400,5000000],
300 |                           [1,-1,2,346000,8777,8777,3401,5060000]])
301 | 
302 |     print("Message array processing",ob.process_order_array(state,msg_array))
303 | 
304 |     print("Process array of messages and return N_levels of L2 state",ob.process_orders_array_l2(state,msgs_array,10))
305 | 
306 |     print("Volume at price test:",ob.get_volume_at_price(state,1,346000,True))
307 | 
308 |     print("Next executable order test:", ob.get_next_executable_order(state,1))
309 | 
310 |     print("Best bid price:", ob.get_best_price(state,1))
311 | 
312 |     print("Best ask price:", ob.get_best_price(state,0))
313 | 
314 |     print("Best bid and ask prices + quants:", ob.get_best_bid_and_ask_inclQuants(state))
315 | 
316 | 


--------------------------------------------------------------------------------
/gymnax_exchange/jaxlobster/data_loading.py:
--------------------------------------------------------------------------------
  1 | from pathlib import Path
  2 | from typing import Generator, Optional
  3 | import pandas as pd
  4 | from itertools import groupby
  5 | import re
  6 | from tqdm import tqdm
  7 | import math
  8 | import numpy as np
  9 | import torch.distributed as dist
 10 | 
 11 | 
 12 | from constants import MESSAGE_TOKEN_DTYPE_MAP, MESSAGE_TOKEN_TYPES, TIME_COL, get_orderbook_token_types
 13 | 
 14 | 
 15 | def _batch(iterable: list[str], n=1) -> Generator[list[str], None, None]:
 16 |     l = len(iterable)
 17 |     for ndx in range(0, l, n):
 18 |         yield iterable[ndx:min(ndx + n, l)]
 19 | 
 20 | def _df_to_str(df: pd.DataFrame, n_msgs: int = 1) -> list[str]:
 21 |     columns = df.columns.tolist()
 22 |     values = df.to_numpy()
 23 |     row_strings = [','.join([f"{col},{val}" for col, val in zip(columns, row) if not pd.isna(val)]) for row in values]
 24 |     concatted_strings = ['\n'.join(rows) for rows in _batch(row_strings, n_msgs)]
 25 |     # if len(concatted_strings[-1].split("\n")) % n_msgs:
 26 |         # TODO: Currently the last incomplete batch is thrown away.
 27 |         # This can be improved by filling up the batch with the next file batch, but it's quiet tedious
 28 |         # concatted_strings = concatted_strings[:-1]
 29 |     return concatted_strings
 30 | 
 31 | def load_message_df(file: str, cast_dtypes: bool = False, nrows: Optional[int] = None) -> pd.DataFrame:
 32 |     df = pd.read_csv(file, dtype=object, header=None, nrows=nrows)
 33 |     removed_indices = []
 34 |     if df.iloc[:, -1].isna().sum():
 35 |         mask = df.iloc[:, -1].isna()
 36 |         removed_indices = df[~mask].index.to_list()
 37 |         df = df[mask]
 38 |     df = df.dropna(axis=1).reset_index(drop=True)
 39 |     try:
 40 |         df.columns = MESSAGE_TOKEN_TYPES
 41 |     except ValueError as e:
 42 |         print(f"Warning: message file {file} was expected to have {len(MESSAGE_TOKEN_TYPES)} columns" + 
 43 |                 f" but got {len(df.columns)}). This file will be skipped.")
 44 |         return pd.DataFrame()
 45 |     if cast_dtypes:
 46 |         df = df.astype(MESSAGE_TOKEN_DTYPE_MAP)
 47 |     return df, removed_indices
 48 | 
 49 | def load_orderbook_df(file: str, cast_dtypes: bool = False) -> pd.DataFrame:
 50 |     df = pd.read_csv(file, dtype=object, header=None)
 51 |     levels = len(df.columns) // 4
 52 |     columns = get_orderbook_token_types(levels)
 53 |     try:
 54 |         df.columns = columns
 55 |     except ValueError as e:
 56 |         print(f"Warning: orderbook file {file} was expected to have {len(columns)} columns," + 
 57 |                 f" but got {len(df.columns)}). This file will be skipped.")
 58 |         return pd.DataFrame()
 59 |     if cast_dtypes:
 60 |         df = df.astype(int)
 61 |     return df
 62 | 
 63 | def extract_date(path: Path) -> Optional[str]:
 64 |     # Use a regex to find the date in the format YYYY-MM-DD
 65 |     match = re.search(r'\d{4}-\d{2}-\d{2}', path.name)
 66 |     return match.group(0) if match else None
 67 | 
 68 | def convert_to_nanoseconds(s: pd.Series):
 69 |     split_times = s.str.split('.', expand=True)
 70 |     split_times.columns = ['seconds', 'fractional']
 71 |     seconds = split_times['seconds'].astype(int)
 72 |     fractional = split_times['fractional'].fillna('0').str.ljust(9, '0').str[:9].astype(int)
 73 |     return seconds * 1_000_000_000 + fractional
 74 | 
 75 | def merge_dfm_dfo(dfm: pd.DataFrame, dfo: pd.DataFrame, n_msgs: int = 50):
 76 |     # reduce n_msgs by 1 due to the prepended orderbook row so that in total one batch has again n_msgs elements
 77 |     n_msgs -= 1
 78 |     num_blocks = math.ceil(len(dfm) / n_msgs)
 79 |     merged_len = len(dfm) + num_blocks
 80 |     merged_df = pd.DataFrame(index=range(merged_len), columns=[*dfm.columns, *dfo.columns])
 81 |     b_indices = np.arange(num_blocks) * (n_msgs + 1)
 82 |     b_vals = np.arange(num_blocks) * n_msgs
 83 |     merged_df.loc[b_indices, dfo.columns] = dfo.iloc[b_vals].values
 84 |     a_indices = np.setdiff1d(np.arange(merged_len), b_indices)
 85 |     merged_df.loc[a_indices, dfm.columns] = dfm.values
 86 |     return merged_df
 87 | 
 88 | def compute_df_from_file_group(group: list[str], n_msgs: int = 50, only_use_message_orderbook_matches: bool = True, differentiate_time: bool = True):
 89 |     df, dfm, dfo = pd.DataFrame(), pd.DataFrame(), pd.DataFrame()
 90 |     if only_use_message_orderbook_matches and len(group) != 2:
 91 |         return pd.DataFrame()
 92 |     for file in list(group):
 93 |         file = str(file)
 94 |         if "message" in file:
 95 |             dfm, removed_indices = load_message_df(file)
 96 |             if dfm.empty:
 97 |                 continue
 98 |             if differentiate_time:
 99 |                 dfm[TIME_COL] = convert_to_nanoseconds(dfm[TIME_COL])
100 |                 dfm[TIME_COL] = dfm[TIME_COL].diff()
101 |                 dfm = dfm.dropna(subset=TIME_COL).reset_index(drop=True)
102 |                 dfm[TIME_COL] = dfm[TIME_COL].astype(int)
103 |             df = dfm.copy()
104 |         elif "orderbook" in file:
105 |             dfo = load_orderbook_df(file)
106 |             if dfo.empty:
107 |                 continue
108 |             df = dfo.copy()
109 |         else:
110 |             print(f"File {file} not known. Expected 'orderbook' or 'message' in file")
111 |     if not dfm.empty and not dfo.empty:
112 |         if removed_indices:
113 |             dfo = dfo.drop(index=removed_indices)
114 |         if differentiate_time:
115 |             dfo = dfo[1:] # adjust to fit to message dataframe
116 |         dfo = dfo.reset_index(drop=True)
117 |         df = merge_dfm_dfo(dfm, dfo, n_msgs)
118 |     return df
119 |     
120 | def get_data_stream_generator(data_dir: str, filter_str: str = "", n_msgs: int = 50, only_use_message_orderbook_matches: bool = True, 
121 |                               differentiate_time: bool = True, world_size: int = 1, rank: int = 0, n_files: int = -1) -> Generator[list[str], None, None]:
122 |     """Generate batches of message files and orderbook files in concatted format (depending on the filter_str) and yield them
123 | 
124 |     :param data_dir: directory of csv files. Expects files to be of format GOOG_2022-01-11_34200000_57600000_message_10.csv
125 |     :param batch_size: batch size that is yielded. If None, return single elements, otherwise return 
126 |         list of length batch size where each element contains n_msgs rows, defaults to None
127 |     :param filter_str: optional filter of files, if e.g. only message files should be processed, defaults to ""
128 |     :param n_msgs: number of messages in one batch element. One message/ row will look like this:
129 | 
130 |         '
131 |         \<time\>,123345567768,\<event_type\>,1,\<order_id\>,123456,\<size\>,100,\<price\>,200,\<direction\>,1,
132 |         \<ask_price_1\>,200,\<ask_size_1\>,20,\<bid_price_1\>,180,\<bid_size_1\>30,\<ask_price_N\>,210,\<ask_size_N\>,20,\<bid_price_N\>,170,\<bid_size_N\>30,
133 |         '
134 |         
135 |         defaults to 20
136 |     :param only_use_message_orderbook_matches: Whether to only use days where message data and orderbook data is available, defaults to True
137 |     :yield: batch or single element of n_msgs
138 |     """
139 |     if filter_str and only_use_message_orderbook_matches:
140 |         print("Warning: Data files are filtered but 'only_use_message_orderbook_matches' is True," +
141 |               " meaning orderbook files and message files are expected. If you want to train purely on one of the two," + 
142 |               " set 'only_use_message_orderbook_matches' to False")
143 |     filter_str = f"*{filter_str}*.csv" if filter_str else "*.csv"
144 | 
145 |     # TODO: There is probably a more efficient way to do this, since extract_date is called two times.
146 |     # But the array is so small that is basically doesn't matter
147 |     msg_files = sorted(list(Path(data_dir).glob(filter_str)), key=extract_date)
148 |     file_groups = [(key, list(group)) for key, group in groupby(msg_files, key=extract_date)]
149 |     # taken_over_batch = None
150 |     # if "test" not in data_dir:
151 |     #     file_groups = file_groups[26:]
152 |     print(f"Unique days: {len(file_groups)}")
153 |     if n_files > 0 and n_files < len(file_groups):
154 |         print(f"Using first {n_files} of {len(file_groups)}")
155 |         file_groups = file_groups[:n_files]
156 |     for _, group in tqdm(file_groups):
157 |         print("Loading file group " + str([str(g) for g in group]))
158 |         df = compute_df_from_file_group(group, n_msgs, only_use_message_orderbook_matches, differentiate_time)
159 |         if df.empty:
160 |             continue
161 |         dfs = _df_to_str(df, n_msgs)
162 |         # if batch_size is None:
163 |         if "test" in data_dir:  # TODO: FIXME: Remove me possibly
164 |             dfs = dfs[:50]  # take only first 50 elements of n_msgs just so validation doesn't take hours
165 |         for i, el in enumerate(dfs):
166 |             if i % world_size == rank:
167 |                 yield el
168 |         # else:
169 |         #     if taken_over_batch:
170 |         #         remainder = batch_size - len(batch)
171 |         #         batch = taken_over_batch + dfs[:remainder]
172 |         #         dfs = dfs[remainder:]
173 |         #         taken_over_batch = None
174 |         #         yield batch
175 |         #     for batch in _batch(dfs, batch_size):
176 |         #         if len(batch) == batch_size:
177 |         #             yield batch
178 |         #         else:
179 |         #             taken_over_batch = batch
180 | 
181 | 
182 | def get_data_stream(data_dir: str, filter_str: str = "", n_msgs: int = 50, only_use_message_orderbook_matches: bool = True, 
183 |                     differentiate_time: bool = True, n_files: int = -1) -> list[str]:
184 |     """Generate batches of message files and orderbook files in concatted format (depending on the filter_str) and yield them
185 | 
186 |     :param data_dir: directory of csv files. Expects files to be of format GOOG_2022-01-11_34200000_57600000_message_10.csv
187 |     :param batch_size: batch size that is yielded. If None, return single elements, otherwise return 
188 |         list of length batch size where each element contains n_msgs rows, defaults to None
189 |     :param filter_str: optional filter of files, if e.g. only message files should be processed, defaults to ""
190 |     :param n_msgs: number of messages in one batch element. One message/ row will look like this:
191 | 
192 |         '
193 |         \<time\>,123345567768,\<event_type\>,1,\<order_id\>,123456,\<size\>,100,\<price\>,200,\<direction\>,1,
194 |         \<ask_price_1\>,200,\<ask_size_1\>,20,\<bid_price_1\>,180,\<bid_size_1\>30,\<ask_price_N\>,210,\<ask_size_N\>,20,\<bid_price_N\>,170,\<bid_size_N\>30,
195 |         '
196 |         
197 |         defaults to 20
198 |     :param only_use_message_orderbook_matches: Whether to only use days where message data and orderbook data is available, defaults to True
199 |     :yield: batch or single element of n_msgs
200 |     """
201 |     if filter_str and only_use_message_orderbook_matches:
202 |         print("Warning: Data files are filtered but 'only_use_message_orderbook_matches' is True," +
203 |               " meaning orderbook files and message files are expected. If you want to train purely on one of the two," + 
204 |               " set 'only_use_message_orderbook_matches' to False")
205 |     filter_str = f"*{filter_str}*.csv" if filter_str else "*.csv"
206 | 
207 |     # TODO: There is probably a more efficient way to do this, since extract_date is called two times.
208 |     # But the array is so small that is basically doesn't matter
209 |     msg_files = sorted(list(Path(data_dir).glob(filter_str)), key=extract_date)
210 |     file_groups = [(key, list(group)) for key, group in groupby(msg_files, key=extract_date)]
211 |     dfs = []
212 |     print(f"Unique days: {len(file_groups)}")
213 |     if n_files > 0 and n_files < len(file_groups):
214 |         print(f"Using first {n_files} file_groups out of {len(file_groups)}")
215 |         file_groups = file_groups[:n_files]
216 |     for _, group in tqdm(file_groups):
217 |         print("Loading file group " + str([str(g) for g in group]))
218 |         df = compute_df_from_file_group(group, n_msgs, only_use_message_orderbook_matches, differentiate_time)
219 |         if df.empty:
220 |             continue
221 |         dfs.extend(_df_to_str(df, n_msgs))
222 |     return dfs


--------------------------------------------------------------------------------
/gymnax_exchange/utils/utils.py:
--------------------------------------------------------------------------------
  1 | import jax
  2 | import jax.numpy as jnp
  3 | import jax.tree_util as jtu
  4 | import random
  5 | import gymnax_exchange.jaxob.JaxOrderBookArrays as job
  6 | """hamilton_apportionment_permuted_jax: A utility function using JAX, 
  7 |                                      implementing a Hamilton apportionment 
  8 |                                      method with randomized seat allocation."""
  9 | 
 10 | 
 11 | def argsort_rev(arr):
 12 |     """ 'arr' sorted in descending order (LTR priority tie-breaker) """
 13 |     return (arr.shape[0] - 1 - jnp.argsort(arr[::-1]))[::-1]
 14 | 
 15 | def rank_rev(arr):
 16 |     """ Rank array in descending order, with ties having left-to-right priority. """
 17 |     return jnp.argsort(argsort_rev(arr))
 18 | 
 19 | @jax.jit
 20 | def clip_by_sum_int(a: jax.Array, max_sum: int) -> jax.Array:
 21 |     """ Clip a vector so that its sum is at most max_sum as an integer,
 22 |         while preserving the relative proportions of the elements.
 23 |         Ties have left-to-right priority.
 24 | 
 25 |         ex: clip_by_sum_int(jnp.array([3, 2, 3, 1]), 8)) -->  [3 2 2 1]
 26 | 
 27 |     Args:
 28 |         a: The vector to clip.
 29 |         max_sum: The maximum sum of the vector.
 30 | 
 31 |     Returns:
 32 |         The clipped vector.
 33 |     """
 34 |     def clip(a: jax.Array, a_sum: int) -> jax.Array:
 35 |         a, remainders = jnp.divmod(a * max_sum, a_sum)
 36 |         rest = max_sum - jnp.sum(a)
 37 |         ranks = rank_rev(remainders)
 38 |         
 39 |         # add 1 to first 'rest' elements of original 'a' with highest remainder
 40 |         a = jnp.where(
 41 |             ranks < rest,
 42 |             a + 1,
 43 |             a,
 44 |         )
 45 |         return a
 46 | 
 47 |     a_sum = jnp.sum(a)
 48 |     return jax.lax.cond(
 49 |         a_sum > max_sum,
 50 |         lambda: clip(a, a_sum),
 51 |         lambda: a,
 52 |     )
 53 | from functools import partial
 54 | @partial(jax.vmap, in_axes=(0, None))
 55 | def p_in_cnl(p, prices_cnl):
 56 |     return jnp.where((prices_cnl == p) & (p != 0), True, False)
 57 | def matching_masks(prices_a, prices_cnl):
 58 |     res = p_in_cnl(prices_a, prices_cnl)
 59 |     return jnp.any(res, axis=1), jnp.any(res, axis=0)
 60 | 
 61 | 
 62 | def tree_stack(trees):
 63 |     return jtu.tree_map(lambda *v: jnp.stack(v), *trees)
 64 | 
 65 | def tree_unstack(tree):
 66 |     leaves, treedef = jtu.tree_flatten(tree)
 67 |     return [treedef.unflatten(leaf) for leaf in zip(*leaves, strict=True)]
 68 | 
 69 | def array_index(array,index):
 70 |     return array[index]
 71 | 
 72 | @jax.jit
 73 | def index_tree(tree,index):
 74 |     array_index = lambda array,index : array[index]
 75 |     indeces=[index]*len(jtu.tree_flatten(tree)[0])
 76 |     tree_indeces=jtu.tree_unflatten(jtu.tree_flatten(tree)[1],indeces)
 77 |     return jtu.tree_map(array_index,tree,tree_indeces)
 78 | 
 79 | def hamilton_apportionment_permuted_jax(votes, seats, key):
 80 |     """
 81 |     Compute the Hamilton apportionment method with permutation using JAX.
 82 | 
 83 |     Args:
 84 |         votes (jax.Array): Array of votes for each party/entity.
 85 |         seats (int): Total number of seats to be apportioned.
 86 |         key (chex.PRNGKey): JAX key for random number generation.
 87 | 
 88 |     Returns:
 89 |         jax.Array: Array of allocated seats to each party/entity.
 90 |     """
 91 |     std_divisor = jnp.sum(votes) / seats # Calculate the standard divisor.
 92 |     # Initial allocation of seats based on the standard divisor and compute remainders.
 93 |     init_seats, remainders = jnp.divmod(votes, std_divisor)
 94 |     # Compute the number of remaining seats to be allocated.
 95 |     remaining_seats = jnp.array(seats - init_seats.sum(), dtype=jnp.int32) 
 96 |     # Define the scanning function for iterative seat allocation.
 97 |     def allocate_remaining_seats(carry,x): # only iterate 4 times, as remaining_seats in {0,1,2,3}
 98 |         key,init_seats,remainders = carry
 99 |         key, subkey = jax.random.split(key)
100 |         # Create a probability distribution based on the maximum remainder.
101 |         distribution = (remainders == remainders.max())/(remainders == remainders.max()).sum()
102 |         # Randomly choose a party/entity to allocate a seat based on the distribution.
103 |         chosen_index = jax.random.choice(subkey, remainders.size, p=distribution)
104 |         # Update the initial seats and remainders for the chosen party/entity.
105 |         updated_init_seats = init_seats.at[chosen_index].add(jnp.where(x < remaining_seats, 1, 0))
106 |         updated_remainders = remainders.at[chosen_index].set(0)
107 |         return (key, updated_init_seats, updated_remainders), x
108 |         # Iterate over parties/entities to allocate the remaining seats.
109 |     (key, init_seats, remainders), _ = jax.lax.scan(
110 |                                                     allocate_remaining_seats,
111 |                                                     (key, init_seats, remainders), 
112 |                                                     xs=jnp.arange(votes.shape[0])
113 |                                                     )
114 |     return init_seats
115 | 
116 | 
117 | def create_init_book(cfg:job.JAXLOB_Configuration,
118 |                      order_capacity=10,
119 |                      trade_capacity=10,
120 |                      pricerange=[2190000,2200000,2210000],
121 |                      quantrange=[0,500],
122 |                      timeinit=[34200,0],
123 |                      percent_fill=0.5):
124 |     """
125 |     Generates a random orderbook state for a given maximum capacity for orders and trades. 
126 |     Random prices/quantities generated by uniform sampling in pricerange/quantrange. 
127 |     """
128 |     qtofill = int(order_capacity*percent_fill) #fill available space.
129 |     asks=[]
130 |     bids=[]
131 |     orderid=cfg.init_id
132 |     traderid=cfg.init_id
133 |     times=timeinit[0]
134 |     timens=timeinit[1]
135 |     for i in range(qtofill):
136 |         asks.append([random.randint(pricerange[1],pricerange[2]),random.randint(quantrange[0],quantrange[1]),orderid,traderid,times,timens])
137 |         bids.append([random.randint(pricerange[0],pricerange[1]),random.randint(quantrange[0],quantrange[1]),orderid-1,traderid-1,times,timens])
138 |         orderid-=2
139 |         traderid-=2
140 |     bids=jnp.concatenate((jnp.array(bids),
141 |                           jnp.ones((order_capacity-qtofill,job.cst.ORDERBOOK_FEAT),dtype=jnp.int32)*job.cst.EMPTY_SLOT),
142 |                           axis=0)
143 |     asks=jnp.concatenate((jnp.array(asks),
144 |                           jnp.ones((order_capacity-qtofill,job.cst.ORDERBOOK_FEAT),dtype=jnp.int32)*job.cst.EMPTY_SLOT),
145 |                           axis=0)
146 |     trades=jnp.ones((trade_capacity,job.cst.TRADE_FEAT),dtype=jnp.int32)*job.cst.EMPTY_SLOT
147 |     return asks,bids,trades
148 | 
149 | def create_rand_message(type='limit',
150 |                         side=None,
151 |                         price_range=[2100000,2200000],
152 |                         quant_range=[0,500],
153 |                         prev_time=job.cst.STARTOFDAY,
154 |                         times_range=[0,2],
155 |                         timens_range=[0,job.cst.NS_PER_SEC]):
156 | 
157 |     type_options = ['limit', 'cancel', 'market']
158 |     side_options = ['bid', 'ask']
159 |     if type == None:
160 |         type = type = random.choice(type_options)
161 |     if side == None:
162 |         side = random.choice(side_options)
163 |     price = random.randint(price_range[0], price_range[1])
164 |     quant = random.randint(quant_range[0], quant_range[1])
165 |     delta_times = random.randint(times_range[0], times_range[1])
166 |     delta_timens = random.randint(timens_range[0], timens_range[1])
167 |     
168 |     
169 |     return create_message(type=type,
170 |                             side=side,
171 |                             price=price,
172 |                             quant=quant,
173 |                             times=prev_time[0]+delta_times,
174 |                             timens=prev_time[1]+delta_timens)
175 | 
176 | def create_message(type='limit',side='bid',price=2200000,quant=10,times=36000,timens=0,id=8888):
177 |     """
178 |     Generates a specific message (based on human-editable inputs)
179 |     Outputs both the 'dictionary' format 
180 |     """
181 |     if type=='limit':
182 |         type_num=job.cst.MessageType.LIMIT.value
183 |     elif type =='cancel' or type == 'delete':
184 |         type_num=job.cst.MessageType.CANCEL.value
185 |     elif type =='market':
186 |         type_num=job.cst.MessageType.MATCH.value
187 |     else:
188 |         raise ValueError('Type is none of: limit, cancel, delete or market')
189 | 
190 |     if side=='bid':
191 |         side_num=job.cst.BidAskSide.BID.value
192 |     elif side =='ask':
193 |         side_num=job.cst.BidAskSide.ASK.value
194 |     else:
195 |         raise ValueError('Side is none of: bid or ask')
196 |     
197 |     dict_msg={
198 |         'side':side_num,
199 |         'type':type_num,
200 |         'price':price,
201 |         'quantity':quant,
202 |         'orderid':id,
203 |         'traderid':id,
204 |         'time':times,
205 |         'time_ns':timens}
206 |     array_msg=jnp.array([type_num,side_num,quant,price,id,id,times,timens])
207 |     return dict_msg,array_msg
208 | 
209 | create_messages=jax.vmap(create_message,in_axes=(None,None,0,0,0,0,0),out_axes=(0,0))
210 | 
211 | 
212 | def get_random_order_to_cancel(book_side,
213 |                                side='bid',
214 |                                prev_time=job.cst.TEST_TIME,
215 |                                times_range=[0,2],
216 |                                timens_range=[0,job.cst.NS_PER_SEC]):
217 |         """
218 |         Obtains a random order ID from the given book side (asks or bids) and returns the order to cancel.
219 |         """
220 |         # Filter out empty slots
221 |         valid_orders = book_side[book_side[:, job.cst.OrderSideFeat.OID.value] != job.cst.EMPTY_SLOT]
222 |         
223 |         
224 |         if valid_orders.shape[0] == 0:
225 |             raise ValueError("No valid orders to cancel.")
226 |         
227 | 
228 |         delta_times = random.randint(times_range[0], times_range[1])
229 |         delta_timens = random.randint(timens_range[0], timens_range[1])
230 |         
231 |         # Select a random order
232 |         random_index = random.randint(0, valid_orders.shape[0] - 1)
233 |         order_to_cancel = valid_orders[random_index]
234 |         
235 |         return create_message(type='cancel',
236 |                               side=side,
237 |                               price=order_to_cancel[job.cst.OrderSideFeat.P.value],
238 |                               quant=order_to_cancel[job.cst.OrderSideFeat.Q.value],
239 |                               times=prev_time[0]+delta_times,
240 |                               timens=prev_time[1]+delta_timens,
241 |                               id=order_to_cancel[job.cst.OrderSideFeat.OID.value])
242 | 
243 | def get_random_aggressive_order(book_side,
244 |                                 side='bid',
245 |                                 prev_time=job.cst.TEST_TIME,
246 |                                 times_range=[0,2],
247 |                                 timens_range=[0,job.cst.NS_PER_SEC]):
248 |         """
249 |         Obtains a price that will guarantee to match, randomly selects a quantity to match which is between 0 and 2x the quantity of the order to match.
250 |         """
251 |         
252 |         if side=='bid':
253 |             best_price=job.get_best_bid(job.JAXLOB_Configuration(),book_side)
254 |             best_vol=job.get_volume_at_price(book_side,best_price)
255 |         else:
256 |             best_price=job.get_best_ask(job.JAXLOB_Configuration(),book_side)
257 |             best_vol=job.get_volume_at_price(book_side,best_price)
258 |         
259 |         delta_times = random.randint(times_range[0], times_range[1])
260 |         delta_timens = random.randint(timens_range[0], timens_range[1])
261 | 
262 |         opp_side= 'ask' if side=='bid' else 'bid'
263 |         
264 |         return create_message(type='limit',
265 |                               side=opp_side,
266 |                               price=best_price,
267 |                               quant= random.randint(0,2*best_vol),
268 |                               times=prev_time[0]+delta_times,
269 |                               timens=prev_time[1]+delta_timens,
270 |                               id=job.cst.DUMMYID)
271 | 
272 | def create_message_forvmap(type='limit',side='bid',price=2200000,quant=10,times=36000,timens=0,nvmap=10):
273 |     if type=='limit':
274 |         type_num=1
275 |     elif type =='cancel' or type == 'delete':
276 |         type_num=2
277 |     elif type =='market':
278 |         type_num=4
279 |     else:
280 |         raise ValueError('Type is none of: limit, cancel, delete or market')
281 | 
282 |     if side=='bid':
283 |         side_num=1
284 |     elif side =='ask':
285 |         side_num=-1
286 |     else:
287 |         raise ValueError('Side is none of: bid or ask')
288 |     
289 |     dict_msg={
290 |     'side':jnp.array([side_num]*nvmap),
291 |     'type':jnp.array([type_num]*nvmap),
292 |     'price':jnp.array([price]*nvmap),
293 |     'quantity':jnp.array([quant]*nvmap),
294 |     'orderid':jnp.array([8888]*nvmap),
295 |     'traderid':jnp.array([8888]*nvmap),
296 |     'time':jnp.array([times]*nvmap),
297 |     'time_ns':jnp.array([timens]*nvmap)}
298 |     array_msg=jnp.array([type_num,side_num,quant,price,8888,8888,times,timens]*nvmap)
299 |     return dict_msg,array_msg
300 | 
301 | 
302 | if __name__ == "__main__":
303 |     # Example configuration object
304 | 
305 |     cfg = job.JAXLOB_Configuration()
306 | 
307 |     key=jax.random.PRNGKey(42)
308 |     key,subkey=jax.random.split(key)
309 |     # Create initial order book
310 |     asks, bids, trades = create_init_book(cfg)
311 | 
312 |     # Print the initial order book
313 |     print("Initial Asks:\n", asks)
314 |     print("Initial Bids:\n", bids)
315 |     print("Initial Trades:\n", trades)
316 | 
317 |     print("Random order to cancel",get_random_aggressive_order(bids))
318 | 
319 |     # # Create a specific message
320 |     # dict_msg, array_msg = create_message()
321 |     # print("\nSpecific Message (Dict):\n", dict_msg)
322 |     # print("Specific Message (Array):\n", array_msg)
323 | 
324 |     # # Create messages for vmap
325 |     # dict_msg_vmap, array_msg_vmap = create_message_forvmap()
326 |     # print("\nMessages for vmap (Dict):\n", dict_msg_vmap)
327 |     # print("Messages for vmap (Array):\n", array_msg_vmap)


--------------------------------------------------------------------------------
/gymnax_exchange/jaxen/Speed_test.py:
--------------------------------------------------------------------------------
  1 | import os, sys, time, dataclasses
  2 | from typing import Tuple, Optional, Dict
  3 | import sys
  4 | import time
  5 | import dataclasses
  6 | import jax
  7 | from jax import vmap
  8 | import jax.numpy as jnp
  9 | import chex
 10 | from flax import struct
 11 | import jax.tree_util as jtu
 12 | from functools import partial
 13 | from typing import Any
 14 | from gymnax_exchange.utils import utils as util
 15 | import pandas as pd
 16 | from datetime import datetime
 17 | #from typing import List, Tuple
 18 | 
 19 | # for debugging
 20 | jax.config.update('jax_disable_jit', False)
 21 | jax.config.update("jax_traceback_in_locations_limit", -1)
 22 | jax.config.update("jax_log_compiles", False)
 23 | 
 24 | #from gymnax_exchange.jaxen.from_JAXMARL.spaces import Box, MultiDiscrete, Discrete
 25 | 
 26 | 
 27 | 
 28 | from gymnax_exchange.jaxob import JaxOrderBookArrays as job
 29 | from gymnax_exchange.jaxob.jaxob_config import MarketMaking_EnvironmentConfig
 30 | from gymnax_exchange.jaxob.jaxob_config import Execution_EnvironmentConfig
 31 | from gymnax_exchange.jaxob.jaxob_config import MultiAgentConfig
 32 | from gymnax_exchange.jaxob.jaxob_config import World_EnvironmentConfig
 33 | 
 34 | from gymnax_exchange.jaxen.marl_env import MARLEnv
 35 | 
 36 | # ----------------------------------------------
 37 | # New VMAP rollout script + timing statistics
 38 | # ----------------------------------------------
 39 | 
 40 | 
 41 | def main():
 42 |     current_datetime = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
 43 |     output_name = f"timing_results_4000_50_compile_before_save_obs_shuffle_action_messages_{current_datetime}"
 44 |     output_file_name = output_name + ".txt"
 45 |     output_file_path = "/home/myuser/gymnax_exchange/jaxen/Timing_speed/" + output_file_name
 46 |     #with open(output_file_path, "w") as f:
 47 |     #    f.write(f"Running with {10} envs and {10} steps\n")
 48 | 
 49 |     rng = jax.random.PRNGKey(30) # TODO i think this should be changed to the new key function in JAX .key()
 50 |     rng, key_reset, key_policy, key_step = jax.random.split(rng, 4)
 51 | 
 52 | 
 53 |     agent_type_options = [
 54 |         [[1,1], [MarketMaking_EnvironmentConfig(), Execution_EnvironmentConfig()]],
 55 |         [[5,5], [MarketMaking_EnvironmentConfig(), Execution_EnvironmentConfig()]],
 56 |         [[10,10], [MarketMaking_EnvironmentConfig(), Execution_EnvironmentConfig()]],
 57 |         #[[1,], [MarketMaking_EnvironmentConfig()]],
 58 |         #[[1,], [Execution_EnvironmentConfig()]],
 59 |         #[[5,], [MarketMaking_EnvironmentConfig()]],
 60 |         #[[5,], [Execution_EnvironmentConfig()]],
 61 |         #[[10,], [MarketMaking_EnvironmentConfig()]],
 62 |         #[[10,], [Execution_EnvironmentConfig()]]
 63 |     ]
 64 |     n_data_msg_options = [100 , 1]
 65 |     #num_envs_options = [1000, 5000, 8000, 10000]
 66 |    # num_steps_options = [1000, 5000]
 67 | 
 68 |     num_steps_num_envs_options = [
 69 |         #[1000, 5000],
 70 |         #[1000, 6000],
 71 |         #[1000, 7000],
 72 |         #[1000, 8000],
 73 |         [50, 4000],
 74 |         #[32, 4000],
 75 |         #[250, 20000],
 76 |         #[120, 40000],
 77 |         #[500, 10000],
 78 |         #[50, 80000],
 79 |         #[3000, 5000],
 80 |         #[3000, 6000],
 81 |         #[3000, 7000],
 82 |         #[3000, 8000],
 83 |     ]
 84 | 
 85 |     save_obs_rewards = True  # Set to False to not save full trajectory
 86 | 
 87 |     results = []
 88 | 
 89 |     with open(output_file_path, "w") as f: 
 90 |         for i, agent_type_option in enumerate(agent_type_options):
 91 |             for n_data_msg_per_step in n_data_msg_options:
 92 |                 for num_step_num_env_option in num_steps_num_envs_options:
 93 |                             # print(f"Running with {num_envs} envs and {num_steps} steps")
 94 | 
 95 |                             # print("\n" + "="*60)
 96 |                             # print("Starting VMAP timing test loop for MARL")
 97 |                             # print("="*60)
 98 |                             # Create a new world config for this run
 99 | 
100 |                             num_steps = num_step_num_env_option[0]
101 |                             num_envs = num_step_num_env_option[1]
102 | 
103 |                             world_config = World_EnvironmentConfig(
104 |                                 n_data_msg_per_step=n_data_msg_per_step,
105 |                             )
106 | 
107 |                             list_of_agents_configs = agent_type_option[1]
108 |                             number_of_agents_per_type = agent_type_option[0]
109 | 
110 |                             print(number_of_agents_per_type)
111 |                             # Create a new MultiAgentConfig for this run
112 |                             multi_agent_config = MultiAgentConfig()
113 |                             # Set the fields (dataclasses are frozen, so use object.__setattr__)
114 |                             object.__setattr__(multi_agent_config, 'world_config', world_config)
115 |                             object.__setattr__(multi_agent_config, 'list_of_agents_configs', list_of_agents_configs)
116 |                             object.__setattr__(multi_agent_config, 'number_of_agents_per_type', number_of_agents_per_type)
117 | 
118 |                             # Re-instantiate the environment and params for this config
119 |                             env = MARLEnv(key=key_reset, multi_agent_config=multi_agent_config)
120 |                             env_params = env.default_params
121 | 
122 |                             # Write config info to file
123 |                             config_str = (
124 |                                 f"\n{'#'*20}\n"
125 |                                 f"Agents per type: {number_of_agents_per_type}\n"
126 |                                 f"n_data_msg_per_step: {n_data_msg_per_step}\n"
127 |                                 f"Agent config: {list_of_agents_configs}\n"
128 |                                 f"{'#'*20}\n"
129 |                             )
130 |                             print(config_str)
131 |                             f.write(config_str)
132 |                             f.flush()
133 | 
134 |                             status_str = (
135 |                                 f"Running with {num_envs} envs and {num_steps} steps\n"
136 |                                 + "="*60 + "\n"
137 |                             )
138 |                             print(status_str)
139 |                             f.write(status_str)
140 |                             f.flush()
141 | 
142 | 
143 |                             NUM_ENVS   = num_envs         # number of parallel environments
144 |                             NUM_STEPS  = num_steps     # total steps per environment
145 |                             MASTER_KEY = jax.random.PRNGKey(0)
146 | 
147 |                             # -------------------------------------------------
148 |                             # 1) Initial reset of all envs (batched)
149 |                             # -------------------------------------------------
150 |                             master_key, *reset_keys = jax.random.split(MASTER_KEY, NUM_ENVS + 1)
151 |                             batched_reset = jax.vmap(env.reset_env, in_axes=(0, None))
152 | 
153 |                             print("Start reset compilation")
154 |                             #reset_jit = batched_reset.lower(jnp.stack(reset_keys), env_params).compile()
155 |                             _ = batched_reset(jnp.stack(reset_keys), env_params)
156 |                             jax.block_until_ready(_)
157 |                             print("start batched reset")
158 | 
159 |                             reset_start = time.time()
160 |                             obs, state  = batched_reset(jnp.stack(reset_keys), env_params)
161 |                             # force execution to finish before timing
162 |                             jax.block_until_ready(state)
163 |                             reset_time  = time.time() - reset_start
164 | 
165 |                              # ------------------------------------
166 |                             # helper: step one whole batch ON device
167 |                             # ------------------------------------
168 |                             def _batched_step(state_batch, key_batch, env_params):
169 |                                 def _single_step(state, key):
170 |                                     subkeys = jax.random.split(key, len(env.action_spaces))
171 |                                     actions = [
172 |                                         jax.vmap(space.sample)(
173 |                                             jax.random.split(sk, n_agents)
174 |                                         )
175 |                                         for sk, space, n_agents in zip(
176 |                                             subkeys,
177 |                                             env.action_spaces,
178 |                                             env.multi_agent_config.number_of_agents_per_type,
179 |                                         )
180 |                                     ]
181 |                                     return env.step(key, state, actions, env_params)
182 | 
183 |                                 return jax.vmap(_single_step, in_axes=(0, 0))(state_batch, key_batch)
184 | 
185 |                             # ------------------------------------
186 |                             # FULL rollout over NUM_STEPS
187 |                             # ------------------------------------
188 |                             @partial(jax.jit, static_argnums=(3,))          #  NUM_STEPS is static
189 |                             def rollout(state0, rng0, env_params, num_steps):
190 |                                 def body(carry, _):
191 |                                     state_batch, rng = carry
192 |                                     rng, *step_keys = jax.random.split(rng, NUM_ENVS + 1)
193 |                                     _, state_batch, _, _, _ = _batched_step(
194 |                                         state_batch, jnp.stack(step_keys), env_params
195 |                                     )
196 |                                     return (state_batch, rng), None
197 | 
198 |                                 (stateN, _), _ = jax.lax.scan(body, (state0, rng0), None, length=num_steps)
199 |                                 return stateN                          
200 | 
201 | 
202 |                             # -------------------------------------------------
203 |                             # 2) Compile first
204 |                             # -------------------------------------------------
205 |                             print("Start compilation")
206 |                             _ = rollout(state, master_key, env_params, NUM_STEPS)
207 |                             jax.block_until_ready(_)
208 | 
209 |                             # Complete pre-compilation
210 |                             #rollout_jit = rollout.lower(state, master_key, env_params, NUM_STEPS).compile()
211 | 
212 |                             print("Start second rollout and measure runtime")
213 | 
214 |                             start = time.time()
215 |                             final_state = rollout(state, master_key, env_params, NUM_STEPS)
216 |                             jax.block_until_ready(final_state)         
217 |                             rollout_time = time.time() - start
218 | 
219 | 
220 | 
221 |                             # -------------------------------------------------
222 |                             # 4) Timing statistics
223 |                             # -------------------------------------------------
224 |                             total_steps       = NUM_STEPS * NUM_ENVS          # every env took NUM_STEPS steps
225 |                             avg_steps_per_env = NUM_STEPS
226 |                             avg_time_per_step = rollout_time / total_steps
227 |                             avg_steps_per_sec = total_steps / rollout_time
228 | 
229 |                             result_str = (
230 |                             f"\n[4] Timing Results\n"
231 |                             f"{'-' * 60}\n"
232 |                             f"Total Envs:           {NUM_ENVS}\n"
233 |                             f"Reset time:           {reset_time:.4f} seconds\n"
234 |                             f"Rollout (steps) time: {rollout_time:.4f} seconds\n"
235 |                             f"Total steps:          {total_steps}\n"
236 |                             f"Avg steps per env:    {avg_steps_per_env:.2f}\n"
237 |                             f"Avg time per step:    {avg_time_per_step:.6f} seconds\n"
238 |                             f"Avg steps per sec:    {avg_steps_per_sec:.2f}\n"
239 |                             f"{'=' * 60}\n"
240 |                             )
241 | 
242 |                             print(result_str)
243 |                             f.write(result_str)
244 |                             f.flush()
245 | 
246 | 
247 |                             results.append({
248 |                                 "num_envs": NUM_ENVS,
249 |                                 "num_steps": NUM_STEPS,
250 |                                 "agents_per_type": number_of_agents_per_type,
251 |                                 "n_data_msg_per_step": n_data_msg_per_step,
252 |                                 "reset_time": reset_time,
253 |                                 "rollout_time": rollout_time,
254 |                                 "total_steps": total_steps,
255 |                                 "avg_steps_per_env": avg_steps_per_env,
256 |                                 "avg_time_per_step": avg_time_per_step,
257 |                                 "avg_steps_per_sec": avg_steps_per_sec,
258 |                             })
259 | 
260 |                             # print("\n[4] Timing Results")
261 |                             # print("-" * 60)
262 |                             # print(f"Total Envs:           {NUM_ENVS}")
263 |                             # print(f"Reset time:           {reset_time:.4f} seconds")
264 |                             # print(f"Rollout (steps) time: {rollout_time:.4f} seconds")
265 |                             # print(f"Total steps:          {total_steps}")
266 |                             # print(f"Avg steps per env:    {avg_steps_per_env:.2f}")
267 |                             # print(f"Avg time per step:    {avg_time_per_step:.6f} seconds")
268 |                             # print(f"Avg steps per sec:    {avg_steps_per_sec:.2f}")
269 |                             # print("=" * 60)
270 | 
271 |     df = pd.DataFrame(results)
272 |     df.to_csv("/home/myuser/gymnax_exchange/jaxen/Timing_speed/" + output_name + ".csv", index=False)
273 | 
274 | 
275 | if __name__ == "__main__":
276 |     main()


--------------------------------------------------------------------------------
/gymnax_exchange/jaxen/base_env.py:
--------------------------------------------------------------------------------
  1 | """
  2 | Base Environment with variable start time for episodes. 
  3 | 
  4 | University of Oxford
  5 | Corresponding Author: 
  6 | Sascha Frey (sascha.frey@st-hughs.ox.ac.uk)
  7 | Kang Li     (kang.li@keble.ox.ac.uk)
  8 | Peer Nagy   (peer.nagy@reuben.ox.ac.uk)
  9 | Valentin Mohl (valentin.mohl@cs.ox.ac.uk)
 10 | Reuben Leyland (Reuben.leyland@sky.com)
 11 | 
 12 | Module Description
 13 | This module provides the base simulation environment for limit order books 
 14 |  using JAX for high-performance computations. It serves as the foundation for 
 15 |  both single-agent and multi-agent reinforcement learning applications in 
 16 |  financial markets, providing core order book simulation functionality.
 17 | 
 18 | Key Components
 19 | LoadedEnvState:   Dataclass to manage the state of the environment, 
 20 |                   including order book states, trade records, and timing information.
 21 | LoadedEnvParams:  Configuration class for environment parameters, 
 22 |                   including message data, book data, and episode timing.
 23 | BaseLOBEnv: Main environment class inheriting from Gymnax's base environment, 
 24 |             providing methods for environment initialization, 
 25 |             stepping through time steps, and resetting the environment. 
 26 | 
 27 | Functionality Overview
 28 | __init__:           Sets up initial values and paths. Loads data from 
 29 |                     LOBSTER and pre-calculates all initial states for 
 30 |                     reset.
 31 | default_params:     Returns the default environment parameters, 
 32 |                     including the preprocessed message and book data.
 33 | step_env:           Advances the environment by one step. It processes both the
 34 |                     action messages and data messages through the order book, 
 35 |                     updates the state, and determines the reward 
 36 |                     and termination condition.
 37 | reset_env:          Resets the environment to an initial state. 
 38 |                     It selects a new data window, initializes the order book, 
 39 |                     and sets the initial state.
 40 | is_terminal:        Checks whether the current state is terminal, 
 41 |                     based on the elapsed time since the episode's start.
 42 | _get_data_messages: Fetches an array of messages for a given step 
 43 |                     within a data window.
 44 | """
 45 | 
 46 | # from jax import config
 47 | # config.update("jax_enable_x64",True)
 48 | import sys
 49 | from ast import Dict
 50 | from contextlib import nullcontext
 51 | from email import message
 52 | from random import sample
 53 | from re import L
 54 | import jax
 55 | import jax.numpy as jnp
 56 | import numpy as np
 57 | from jax import lax
 58 | from gymnax.environments import environment, spaces
 59 | from typing import Tuple, Optional,Union
 60 | import chex
 61 | from flax import struct
 62 | import itertools
 63 | from gymnax_exchange.jaxob import JaxOrderBookArrays as job
 64 | from gymnax_exchange.jaxlobster.lobster_loader import LoadLOBSTER_resample
 65 | #from gymnax_exchange.jaxlobster.gen_loader import GenLoader
 66 | from gymnax_exchange.utils.utils import *
 67 | import pickle
 68 | from jax.experimental import checkify
 69 | import os
 70 | 
 71 | #Config File:
 72 | from gymnax_exchange.jaxob.jaxob_config import World_EnvironmentConfig
 73 | from gymnax_exchange.jaxen.StatesandParams import LoadedEnvParams, LoadedEnvState, WorldState
 74 | 
 75 | 
 76 | 
 77 | class BaseLOBEnv(environment.Environment):
 78 | 
 79 |     """The basic RL environment for the limit order book (LOB) using
 80 |     JAX-LOB functions for manipulating the orderbook. 
 81 |     Inherits from gymnax base environment. 
 82 |     Attributes
 83 |     ----------
 84 |     window_selector : int
 85 |         -1 to randomly choose start times from all available.
 86 |         int in range(0,n_starts) to choose a specific window for debug. 
 87 |     data_type : str
 88 |         "fixed_steps" and "fixed_time" to defn episode end crit.
 89 |     episode_time : int
 90 |         Length of episode in steps or seconds based on above.
 91 |     n_data_msg_per_step : int
 92 |         number of messages to process per step. 
 93 |     day_start : int
 94 |         Beginning time of day in seconds
 95 |     day_end : int
 96 |         End time of day in seconds
 97 |     nOrdersPerSide : int
 98 |         Maximum capacity of orders for JAXLOB
 99 |     nTradesLogged : int
100 |         Maximum number of trades logged (in a step)
101 |     book_depth : int
102 |         Depth considered for LOBSTER data retrieval
103 |     n_actions : int
104 |         Number of actions. Dimension of act space
105 |     n_ticks_in_book : int
106 |         Depth of passive order in act space in ticks.
107 |     customIDCounter : int
108 |         Ensures unique IDs for orders submitted by agent.
109 |     trader_unique_id : int
110 |         Offset of unique ID that can be used.
111 |     tick_size : int
112 |         Price tick size. Lobster counts in hundreths of cents.
113 |     start_resolution: int
114 |         Interval, in seconds, at which episodes may start.
115 |     loader : LoadLOBSTER 
116 |         Object that deals with data-loading.
117 |     max_messages_in_episode_arr : jnp.Array 
118 |         Total messages for each possible window.
119 |     messages : jnp.Array  
120 |         Loaded message data.
121 |     books : jnp.Array  
122 |         Loaded book data for start-points
123 |     n_windows : int 
124 |         Number of start points
125 |     start_indeces : jnp.Array  
126 |         Ineces for start points
127 |     end_indeces : jnp.Array  
128 |         Indeces for ep end for each start-point. 
129 |     init_states_array : jnp.Array  
130 |         Initial state for each start point: for reset func. 
131 | 
132 |     Methods
133 |     -------
134 |     info(additional=""):
135 |         Prints the person's name and age.
136 |     """
137 |     def __init__(self,cfg:World_EnvironmentConfig, key):
138 |         super().__init__()
139 |         self.window_selector = cfg.window_selector
140 |         self.ep_type = cfg.ep_type # fixed_steps, fixed_time
141 |         self.episode_time = cfg.episode_time # counted by seconds, 1800s=0.5h or steps
142 |         self.n_data_msg_per_step = cfg.n_data_msg_per_step
143 |         self.day_start = cfg.day_start  # 09:30
144 |         self.day_end = cfg.day_end  # 16:00
145 |         self.nOrdersPerSide=cfg.nOrdersPerSide #100
146 |         self.nTradesLogged=cfg.nTradesLogged
147 |         self.book_depth=cfg.book_depth
148 |         self.n_ticks_in_book = cfg.n_ticks_in_book 
149 |         self.customIDCounter=cfg.customIDCounter
150 |         self.tick_size=cfg.tick_size
151 |         self.start_resolution = cfg.start_resolution  # Use value from config
152 |         self.cfg = cfg
153 | 
154 |         loader=LoadLOBSTER_resample(self.cfg.dataPath,
155 |                                     self.cfg.alphatradePath,
156 |                                 self.book_depth,
157 |                                 self.ep_type,
158 |                                 window_length=self.episode_time,
159 |                                 n_data_msg_per_step=self.n_data_msg_per_step,
160 |                                 window_resolution=self.start_resolution,
161 |                                 day_start=self.day_start,
162 |                                 day_end=self.day_end,
163 |                                 stock=self.cfg.stock,
164 |                                 time_period=self.cfg.timePeriod) 
165 |         msgs,starts,ends,books,max_messages_arr=loader.run_loading()
166 | 
167 | 
168 |         self.max_messages_in_episode_arr = max_messages_arr
169 |         self.messages=msgs #Is different to trad. base: all msgs concat. TODO this should not be saved here
170 |         self.books=books
171 |         self.n_windows = starts.shape[0]
172 |         self.start_indeces=starts
173 |         self.end_indeces=ends
174 |         self._init_states(key,self.cfg.alphatradePath,self.start_indeces)
175 | 
176 |     
177 |     @property
178 |     def default_params(self) -> LoadedEnvParams:
179 |         # Default environment parameters
180 |         return LoadedEnvParams(
181 |             message_data=jnp.asarray(self.messages), 
182 |             book_data=jnp.asarray(self.books),
183 |             init_states_array=self.init_states_array
184 |         )
185 |     
186 |     @partial(jax.jit, static_argnums=(0,))
187 |     def step_env(
188 |         self, key: chex.PRNGKey, state: LoadedEnvState, action: Dict, params: LoadedEnvParams
189 |     ) -> Tuple[chex.Array, LoadedEnvState, float, bool, dict]:
190 |         #Obtain the messages for the step from the message data
191 |         data_messages=self._get_data_messages(params.message_data,
192 |                                               state.start_index,
193 |                                               state.step_counter,
194 |                                               state.init_time[0]+self.cfg.episode_time)
195 |         #data_messages=self._get_generative_messages(params.message_data,n_messages)
196 |         
197 |         #Note: Action of the base environment should consistently be "DO NOTHING"
198 | 
199 |         total_messages=data_messages
200 | 
201 |         #Save time of final message to add to state
202 |         time=total_messages[-1:][0][-2:]
203 |         #Process messages of step (action+data) through the orderbook
204 |         
205 |         ordersides=job.scan_through_entire_array(self.cfg,key,total_messages,(state.ask_raw_orders,state.bid_raw_orders,state.trades))
206 | 
207 |         #Update state (ask,bid,trades,init_time,current_time,OrderID counter,window index for ep, step counter)
208 |         state = LoadedEnvState(ordersides[0],ordersides[1],ordersides[2],state.init_time,\
209 |             state.window_index,state.step_counter+1,state.max_steps_in_episode,state.start_index)
210 |         done = self._internal_terminal_debug(state,params,time)
211 |         reward=0
212 |         #jax.debug.print("Final state after step: \n {}", state)
213 |         return self._get_obs(state,params),state,reward,done,{"info":0}
214 | 
215 |     def reset_env(
216 |         self, key: chex.PRNGKey, params: LoadedEnvParams
217 |     ) -> Tuple[int,LoadedEnvState]:
218 |         """Reset environment state by sampling initial position in OB."""
219 |         idx_data_window = jnp.where(
220 |             self.cfg.window_selector == -1,
221 |             jax.random.randint(key, minval=0, maxval=self.n_windows, shape=()),  
222 |             jnp.array(self.cfg.window_selector, dtype=jnp.int32))
223 |         #jax.debug.print("idx_data_window: {}", idx_data_window)
224 |         first_state = index_tree(params.init_states_array, idx_data_window)
225 |         # def debug_callback(first_state,idx_data_window):
226 |         #     if idx_data_window == 427:  # Debugging for specific window index
227 |         #         print("Debugging reset for window index:", idx_data_window)
228 |         #         print("Resetting environment to initial state for window index:", first_state.window_index)
229 |         #         print("First state details:", first_state)
230 |         # jax.debug.callback(debug_callback, first_state, idx_data_window)
231 |         return 0,first_state
232 |     
233 |     def _internal_terminal_debug(self, state: LoadedEnvState, params: LoadedEnvParams,time : chex.Array) -> bool:
234 |         return (time-state.init_time)[0]>=self.cfg.episode_time
235 | 
236 | 
237 |     # def is_terminal(self, state: LoadedEnvState, params: LoadedEnvParams) -> bool:
238 |     #     """Check whether state is terminal."""
239 |     #     #jax.debug.print("Time: {} , Init time: {}, Difference: {}",state.time, state.init_time,(state.time-state.init_time)[0])
240 |     #     return (state.time-state.init_time)[0]>=self.cfg.episode_time
241 | 
242 |     def _get_state_from_data(self,key,first_message,book_data,max_steps_in_episode,window_index,start_index)->LoadedEnvState:
243 |         time=jnp.array(first_message[-2:])
244 |         #Get initial orders (2xNdepth)x6 based on the initial L2 orderbook for this window 
245 |         def get_initial_orders(book_data,time):
246 |             orderbookLevels=10
247 |             initid=self.cfg.init_id
248 |             #jax.debug.print("\n=== Debug Order Book Initialization ===")
249 |             #jax.debug.print("Raw book_data shape: {}", book_data.shape)
250 |             #jax.debug.print("Raw book_data: {}", book_data)
251 |             data=jnp.array(book_data).reshape(int(10*2),2)
252 |             #jax.debug.print("\nReshaped data: {}", data)
253 |             newarr = jnp.zeros((int(orderbookLevels*2),8),dtype=jnp.int32)
254 |             initOB = newarr \
255 |                 .at[:,3].set(data[:,0]) \
256 |                 .at[:,2].set(data[:,1]) \
257 |                 .at[:,0].set(1) \
258 |                 .at[0:orderbookLevels*4:2,1].set(-1) \
259 |                 .at[1:orderbookLevels*4:2,1].set(1) \
260 |                 .at[:,4].set(initid) \
261 |                 .at[:,5].set(initid-jnp.arange(0,orderbookLevels*2)) \
262 |                 .at[:,6].set(time[0]) \
263 |                 .at[:,7].set(time[1])
264 |             #jax.debug.print("\nFinal initOB array:")
265 |             #jax.debug.print("Shape: {}", initOB.shape)
266 |             #jax.debug.print("Content: {}", initOB)
267 |             #jax.debug.print("Side assignments (column 1): {}", initOB[:,1])
268 |             #jax.debug.print("Price assignments (column 3): {}", initOB[:,3])
269 |             #jax.debug.print("Quantity assignments (column 2): {}", initOB[:,2])
270 |             return initOB
271 |         init_orders=get_initial_orders(book_data,time)
272 |         #jax.debug.print("init_orders {}",init_orders)
273 |         #Initialise both sides of the book as being empty
274 |         asks_raw=job.init_orderside(self.cfg.nOrders)
275 |         bids_raw=job.init_orderside(self.cfg.nOrders)
276 |         trades_init=(jnp.ones((self.cfg.nTrades,8))*-1).astype(jnp.int32)
277 |         #Process the initial messages through the orderbook
278 |         ordersides=job.scan_through_entire_array(self.cfg,key,init_orders,(asks_raw,bids_raw,trades_init))
279 |         
280 |  
281 |         
282 |         return LoadedEnvState(ask_raw_orders=ordersides[0],
283 |                         bid_raw_orders=ordersides[1],
284 |                         trades=ordersides[2],
285 |                         init_time=jnp.array([(window_index*self.start_resolution) 
286 |                                                         %(self.day_end-self.day_start-self.episode_time+self.start_resolution)
287 |                                                         +self.day_start,0])
288 |                                     if self.ep_type=="fixed_time" else time,
289 |                         window_index=window_index,
290 |                         max_steps_in_episode=max_steps_in_episode,
291 |                         start_index=start_index,
292 |                         step_counter=0,
293 | )
294 | 
295 |     def _init_states(self,key,alphatradePath,starts):
296 |         print("START:  pre-reset in the initialization")
297 |         os.makedirs(alphatradePath + '/pre_reset_states/', exist_ok=True)
298 |         pkl_file_name = (alphatradePath + '/pre_reset_states/'
299 |                          + 'ResetState_window_resolution_' + str(self.cfg.start_resolution)
300 |                          + '_eptype_"' + str(self.cfg.ep_type)
301 |                          + '"_depth_' + str(self.cfg.book_depth)
302 |                          + "_stock_" + str(self.cfg.stock)
303 |                          + "_windowidx_"+str(self.cfg.window_selector)
304 |                          + "_nMsgPerStep_"+str(self.cfg.n_data_msg_per_step)
305 |                          + "_episode_time_"+str(self.cfg.episode_time)
306 |                          + "_TimePeriod_"+str(self.cfg.timePeriod)
307 |                          + '.pkl')
308 |         print("pre-reset will be saved to ", pkl_file_name)
309 |         try:
310 |             if self.cfg.use_pickles_for_init:
311 |                 with open(pkl_file_name, 'rb') as f:
312 |                     self.init_states_array = pickle.load(f)
313 |                     print("LOADING STATES FROM PKL...")
314 |             else:
315 |                 raise ValueError("Throw error so re-computes")
316 |         except:
317 |             print("COMPUTING INIT STATES...")
318 |             #for i in range(self.n_windows):
319 |                 #print("message starts",self.messages[starts[i]])
320 |             get_state_jitted= jax.jit(self._get_state_from_data)
321 | 
322 |             states = [get_state_jitted(key,
323 |                                                 self.messages[starts[i]],
324 |                                                 self.books[i],
325 |                                                 self.max_messages_in_episode_arr[i]
326 |                                                     //self.n_data_msg_per_step+1,
327 |                                                     i,
328 |                                                     starts[i]) 
329 |                         for i in range(self.n_windows)]
330 |             self.init_states_array=tree_stack(states)
331 |             print("SAVING STATES TO PKL...")
332 |             with open(pkl_file_name, 'wb') as f:
333 |                 pickle.dump(self.init_states_array, f)
334 |         print("DONE: pre-reset in the initialization")
335 | 
336 |     def _get_obs(self, state: LoadedEnvState, params:LoadedEnvParams) -> chex.Array:
337 |         """Return dummy observation."""
338 |         return 0
339 |     
340 |     def _get_data_messages(self,messageData,start,step_counter,end_time_s):
341 |         """Returns an array of messages for a given step. 
342 |             Parameters:
343 |                     messageData (Array): 2D array of all msgs with
344 |                                         dimensions: messages, features.
345 |                     start (int): Index of first message to in episode
346 |                     step_counter (int): desired step to consider
347 |                     end_time_s (int): End time of ep in seconds
348 |             Returns:
349 |                     Messages (Array): 2D array of messages for step 
350 |         """
351 |         index_offset=start+self.n_data_msg_per_step*step_counter
352 |         
353 |         messages=jax.lax.dynamic_slice_in_dim(messageData,index_offset,self.n_data_msg_per_step,axis=0)
354 |         #jax.debug.print("{}",messages)
355 |         #jax.debug.print("End time: {}",end_time_s)
356 |         #messages=messageData[index_offset:(index_offset+self.n_data_msg_per_step),:]
357 |         #Replace messages after the cutoff time with padded 0s (except time)
358 |         #jax.debug.print("m_wout_time {}",jnp.transpose(jnp.resize(messages[:,-2]>=end_time_s,messages[:,:-2].shape[::-1])))
359 |         if self.cfg.ep_type == "fixed_time":
360 |             # If fixed time, we need to remove messages that are after the end time
361 |             # We do this by replacing the messages with 0s if they are after the end time
362 |             m_wout_time = jnp.where(jnp.transpose(jnp.resize(
363 |                                 messages[:,-2]>=end_time_s,
364 |                                 messages[:,:-2].shape[::-1])),
365 |                               jnp.zeros_like(messages[:,:-2]),
366 |                               messages[:,:-2])
367 |         #jax.debug.print("m_wout_time {}",m_wout_time)
368 | 
369 |             messages=jnp.concatenate((m_wout_time,messages[:,-2:]),axis=1,dtype=jnp.int32)
370 |         return messages
371 |     
372 |     def _get_generative_messages(self,previous_messages,n_messages):
373 |         """Draft Generative Loader:
374 |         Inputs:
375 |         
376 |         Outputs:
377 |         
378 |         
379 |         """
380 |         loader = GenLoader(
381 |         model=GenLoader.dummy_model,
382 |         initial_context=previous_messages,
383 |         initial_ob_state=jnp.zeros((2, 10)),
384 |         n_messages=100
385 |         )
386 |         messages, _ = loader.generate_step()
387 |         
388 | 
389 |         return messages
390 |     
391 |     def _get_pass_price_quant(self, state):
392 |             """Get price and quanitity n_ticks into books"""
393 |             bid_passive_2=state.best_bids[-1, 0] - self.tick_size * self.cfg.n_ticks_in_book
394 |             ask_passive_2=state.best_asks[-1, 0] + self.tick_size * self.cfg.n_ticks_in_book
395 |             quant_bid_passive_2 = job.get_volume_at_price(state.bid_raw_orders, bid_passive_2)
396 |             quant_ask_passive_2 = job.get_volume_at_price(state.ask_raw_orders, ask_passive_2)
397 |             return bid_passive_2,quant_bid_passive_2,ask_passive_2,quant_ask_passive_2
398 |         
399 | 
400 | 
401 |     @property
402 |     def name(self) -> str:
403 |         """Environment name."""
404 |         return "alphatradeBase-v0"
405 | 
406 |     @property
407 |     def num_actions(self) -> int:
408 |         """Number of actions possible in environment."""
409 |         return self.n_actions
410 | 
411 |     def action_space(
412 |         self, params: Optional[LoadedEnvParams] = None
413 |     ) -> Union[spaces.Discrete,spaces.Dict]:
414 |         """Action space of the environment."""
415 |         return spaces.Dict(
416 |             {
417 |                 "sides":spaces.Box(0,2,(self.n_actions,),dtype=jnp.int32),
418 |                 "quantities":spaces.Box(0,100,(self.n_actions,),dtype=jnp.int32),
419 |                 "prices":spaces.Box(0,99999999,(self.n_actions,),dtype=jnp.int32)
420 |             }
421 |         )
422 | 
423 |     # Not used.
424 |     def observation_space(self, params: LoadedEnvParams):
425 |         """Observation space of the environment."""
426 |         return NotImplementedError
427 | 
428 |     # Not used.
429 |     def state_space(self, params: LoadedEnvParams) -> spaces.Dict:
430 |         """State space of the environment.  # Not used.
431 |         """
432 |         return spaces.Dict(
433 |             {
434 |                 "bids": spaces.Box(-1,999999999,shape=(6,self.cfg.nOrders),dtype=jnp.int32),
435 |                 "asks": spaces.Box(-1,999999999,shape=(6,self.cfg.nOrders),dtype=jnp.int32),
436 |                 "trades": spaces.Box(-1,999999999,shape=(8,self.cfg.nTrades),dtype=jnp.int32),
437 |                 "time": spaces.Discrete(params.max_steps_in_episode),
438 |             }
439 |         )
440 |     
441 | 
442 | 
443 |     
444 | if __name__ == "__main__":
445 |     world_config = World_EnvironmentConfig()
446 | 
447 |     rng = jax.random.PRNGKey(0)
448 |     rng,key_init, key_reset, key_policy, key_step = jax.random.split(rng, 5)
449 | 
450 | 
451 |     env= BaseLOBEnv(cfg=world_config,
452 |                     key=key_init)
453 |     
454 |     jax.profiler.start_trace("/tmp/profile-data")
455 |     env_params=env.default_params
456 | 
457 |     obs,state=env.reset(key_reset,env_params)
458 |     # done=False
459 | 
460 |     # while not done :
461 |     obs,state,rewards,done,info=env.step_env(key_step,state,{},env_params)
462 |     rng,key_step = jax.random.split(rng, 2)
463 |     obs,state,rewards,done,info=env.step_env(key_step,state,{},env_params)
464 | 
465 |     rng,key_reset = jax.random.split(rng, 2)
466 |     obs,state=env.reset(key_reset,env_params)
467 | 
468 |     #     print(done)
469 | 
470 |     jax.block_until_ready(state)
471 |     jax.profiler.stop_trace()
472 | 
473 | 
474 | 
475 |     # print(state)
476 |     # print(obs)
477 | 


--------------------------------------------------------------------------------
/gymnax_exchange/jaxrl/MARL/baseline_eval/plotting_episodes.py:
--------------------------------------------------------------------------------
  1 | import os
  2 | from re import L
  3 | import time
  4 | import numpy as np
  5 | import pickle
  6 | import glob
  7 | 
  8 | import matplotlib.pyplot as plt
  9 | 
 10 | os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "0.1"
 11 | 
 12 | 
 13 | from gymnax_exchange.jaxen.marl_env import MARLEnv
 14 | from gymnax_exchange.jaxob.jaxob_config import MultiAgentConfig,Execution_EnvironmentConfig, World_EnvironmentConfig
 15 | from baseline_JAXMARL import Transition
 16 | import argparse
 17 | 
 18 | def get_latest_pickle_file(directory="trajectories", combo_desc="default"):
 19 |     """Find the most recently created pickle file in the specified directory."""
 20 |     if not os.path.exists(directory):
 21 |         raise FileNotFoundError(f"Directory '{directory}' not found")
 22 |     
 23 |     pickle_files = glob.glob(os.path.join(directory, f"*{combo_desc}*.pkl"))
 24 |     if not pickle_files:
 25 |         raise FileNotFoundError(f"No pickle files found in directory '{directory}'")
 26 |     
 27 |     # Sort files by creation time (most recent last)
 28 |     latest_file = max(pickle_files, key=os.path.getctime)
 29 |     return latest_file
 30 | 
 31 | def main():
 32 |     def parse_args():
 33 |         """Parse command line arguments."""
 34 |         parser = argparse.ArgumentParser(description="Plot episode features from trajectory data")
 35 |         parser.add_argument("--directory", "-d", type=str, default="trajectories", 
 36 |                             help="Directory containing trajectory pickle files")
 37 |         parser.add_argument("--combo", "-c", type=str, nargs='+', default=["BB"], 
 38 |                     help="Combo description(s) to filter pickle files (can provide multiple)")
 39 |         parser.add_argument("--save", type=str, default="intra-episode-figs", 
 40 |                     help="Save Directory for plots (can provide multiple)")
 41 |         return parser.parse_args()
 42 | 
 43 |     args = parse_args()
 44 |     """Load trajectory batch from latest pickle file and generate plots."""
 45 | 
 46 |     plot_same_axis([2],["quant_left"],args.combo,input_dir=args.directory, output_dir=args.save+f"/single_plots")
 47 | 
 48 |     # for combo in args.combo:
 49 |     #     try:
 50 |     #         # Find and load the latest pickle file for each combo
 51 |     #         latest_file = get_latest_pickle_file(directory=args.directory, combo_desc=combo)
 52 |     #         print(f"Loading trajectory data from: {latest_file}")
 53 |             
 54 |     #         with open(latest_file, "rb") as f:
 55 |     #             traj_batch = pickle.load(f)
 56 |     #         print(f"Loaded trajectory batch with {len(traj_batch)} agents.")
 57 |     #         # Plot episode features
 58 |     #         # plot_episode_features(traj_batch, output_dir=args.save+f"/{combo}",)
 59 |     #         # plot_specific(traj_batch,[0],["quant_left"], output_dir=args.save+f"/{combo}")
 60 |     #         print(f"Plotting complete. Check the 'intra-episode-figs' directory for output.")
 61 |             
 62 |     #     except Exception as e:
 63 |     #         print(f"Error processing combo '{combo}': {e}")
 64 | 
 65 | def plot_same_axis(env_indices, features,combos,input_dir="", output_dir="intra-episode-figs", feature_names=None, obs_features=None):
 66 |     num_agent_types = 2
 67 |     os.makedirs(output_dir, exist_ok=True)
 68 |     cmap = plt.cm.get_cmap('seismic', num_agent_types)
 69 |     num_envs= len(env_indices)
 70 |     info_fig, info_axes = plt.subplots(num_envs, 1, figsize=(7, 5*num_envs), sharex=True)
 71 |     if num_envs == 1:
 72 |         info_axes = [info_axes] 
 73 |     for c_indx, c in enumerate(combos):
 74 |         latest_file = get_latest_pickle_file(directory=input_dir, combo_desc=c)
 75 |         print(f"Loading trajectory data from: {latest_file}")
 76 |         
 77 |         with open(latest_file, "rb") as f:
 78 |             traj_batch = pickle.load(f)
 79 |         print(f"Loaded trajectory batch with {len(traj_batch)} agents.")
 80 |         for key in features:
 81 | 
 82 |             steps = np.arange(traj_batch[0].action.shape[0])
 83 |             print(info_axes)
 84 | 
 85 |             # Plot this metric for each agent type on the same subplot
 86 |             for env_idx,env_actual in enumerate(env_indices):
 87 |                 for agent_idx, traj in enumerate(traj_batch):
 88 |                     if 'agent' in traj.info and key in traj.info['agent']:
 89 |                         values = traj.info['agent'][key]      
 90 |                         if len(values.shape) <= 2:  # Only plot simple scalar features
 91 |                             env_values = values[:, env_actual] if len(values.shape) > 1 else values
 92 |                             info_axes[env_idx].plot(steps, env_values, color=cmap(c_indx), label=f"Execution Agent - {c}")
 93 |                             info_axes[env_idx].set_title(f"Trajectory plot for {key} measure")
 94 |                             info_axes[env_idx].set_xlabel("Steps")
 95 |                             info_axes[env_idx].set_ylabel(key)
 96 |                             info_axes[env_idx].legend()
 97 |                             info_axes[env_idx].grid(True)
 98 |                         else:
 99 |                             print(f"Skipping plotting for {key} as it has more than 2 dimensions. {values.shape}")
100 |             # Save the figure for this metric
101 |             metric_path = os.path.join(output_dir, f"Mega_plot_{time.strftime('%Y%m%d-%H%M%S')}.png")
102 |             info_fig.tight_layout()
103 |             print(f"Saving plot to: {metric_path}")
104 |             info_fig.savefig(metric_path)
105 |             plt.close(info_fig)
106 | 
107 | 
108 | def plot_specific(traj_batch,env_indices,features, output_dir="intra-episode-figs", feature_names=None, obs_features=None):
109 |     # Create output directory if it doesn't exist
110 |     os.makedirs(output_dir, exist_ok=True)
111 |     
112 |     num_agent_types = len(traj_batch)
113 |     # Create a colormap that can handle all agent types
114 |     cmap = plt.cm.get_cmap('viridis', num_agent_types)
115 |     
116 |     # Get the number of environments
117 |     num_envs = len(env_indices)
118 |     
119 |     # Create figure with subplots for each environment
120 |     fig, axes = plt.subplots(num_envs, 1, figsize=(12, 5*num_envs), sharex=True)
121 |     if num_envs == 1:
122 |         axes = [axes]
123 |     
124 |     # Plot rewards - one subplot per environment on a single figure
125 |     fig_rewards = plt.figure(figsize=(12, 5*num_envs))
126 |     for env_idx in env_indices:
127 |         ax_reward = fig_rewards.add_subplot(num_envs, 1, env_idx+1)
128 |         steps = np.arange(traj_batch[0].reward.shape[0])
129 |         
130 |         # For each agent type, plot a line with different color
131 |         for agent_idx, traj in enumerate(traj_batch):
132 |             rewards = traj.reward
133 |             env_rewards = rewards[:, env_idx] if len(rewards.shape) > 1 else rewards
134 |             ax_reward.plot(steps, env_rewards, color=cmap(agent_idx), label=f"Agent Type {agent_idx}")
135 |         
136 |         ax_reward.set_title(f"Environment {env_idx} Rewards")
137 |         ax_reward.set_xlabel("Steps")
138 |         ax_reward.set_ylabel("Reward")
139 |         ax_reward.legend()
140 |         ax_reward.grid(True)
141 |     
142 |     fig_rewards.tight_layout()
143 |     output_path = os.path.join(output_dir, f"rewards_all_envs_{time.strftime('%Y%m%d-%H%M%S')}.png")
144 |     fig_rewards.savefig(output_path)
145 |     plt.close(fig_rewards)
146 |     
147 |     # Plot actions - one subplot per environment on a single figure
148 |     fig_actions = plt.figure(figsize=(12, 5*num_envs))
149 |     for env_idx in env_indices:
150 |         ax_action = fig_actions.add_subplot(num_envs, 1, env_idx+1)
151 |         steps = np.arange(traj_batch[0].action.shape[0])
152 |         
153 |         # For each agent type, plot a line with different color
154 |         for agent_idx, traj in enumerate(traj_batch):
155 |             actions = traj.action
156 |             env_actions = actions[:, env_idx] if len(actions.shape) > 1 else actions
157 |             ax_action.plot(steps, env_actions, color=cmap(agent_idx), label=f"Agent Type {agent_idx}")
158 |         
159 |         ax_action.set_title(f"Environment {env_idx} Actions")
160 |         ax_action.set_xlabel("Steps")
161 |         ax_action.set_ylabel("Action")
162 |         ax_action.legend()
163 |         ax_action.grid(True)
164 |     
165 |     fig_actions.tight_layout()
166 |     output_path = os.path.join(output_dir, f"actions_all_envs_{time.strftime('%Y%m%d-%H%M%S')}.png")
167 |     fig_actions.savefig(output_path)
168 |     plt.close(fig_actions)
169 |     
170 |     # Plot agent-specific info
171 |     if hasattr(traj_batch[0], 'info') and traj_batch[0].info is not None:
172 |         # Collect all unique keys from agent info across all agent types
173 |         all_keys = set()
174 |         for agent_idx, traj in enumerate(traj_batch):
175 |             if 'agent' in traj.info:
176 |                 agent_info = traj.info['agent']
177 |                 all_keys.update(agent_info.keys())
178 |         print(all_keys)
179 |         
180 |         # For each metric, create a figure with subplots for each environment
181 |         for key in all_keys:
182 |             if key not in features:
183 |                 continue
184 |             for agent_idx, traj in enumerate(traj_batch):
185 |                 if 'agent' in traj.info and key in traj.info['agent']:
186 |                     values = traj.info['agent'][key]
187 |             # If values is a dict, create a subplot for each key in the dict
188 |             if isinstance(values, dict):
189 |                 # Create a figure with subplots for each environment
190 |                 all_sub_keys = set()
191 |                 for agent_idx, traj in enumerate(traj_batch):
192 |                     if 'agent' in traj.info:
193 |                         agent_info = traj.info['agent'][key]
194 |                         all_sub_keys.update(agent_info.keys())
195 |                 print(all_sub_keys)
196 |                 subkey_to_index = {subkey: i for i, subkey in enumerate(all_sub_keys)}
197 |                 info_fig, info_axes = plt.subplots(num_envs, len(all_sub_keys), figsize=(6*len(all_sub_keys), 5*num_envs), sharex=True)
198 |                 if num_envs == 1:
199 |                     info_axes = [info_axes]
200 |                 for env_idx in range(num_envs):
201 |                     for agent_idx, traj in enumerate(traj_batch):
202 |                         if 'agent' in traj.info and key in traj.info['agent']:
203 |                             values = traj.info['agent'][key]      
204 |                             print(values.keys())
205 |                             for subkey in all_sub_keys:
206 |                                 if subkey not in values:
207 |                                     # print(f"Skipping plotting for {key}.{subkey} as it is not present in agent info.")
208 |                                     continue
209 |                                 subvalues = values[subkey]
210 |                                 if len(subvalues.shape) <= 2:  # Only plot simple scalar features
211 |                                     # print(f"Plotting {key}.{subkey} for agent type {agent_idx} in environment {env_idx}")
212 |                                     env_values = subvalues[:, env_idx] if len(subvalues.shape) > 1 else subvalues
213 |                                     info_axes[env_idx,subkey_to_index[subkey]].plot(steps, env_values, color=cmap(agent_idx), 
214 |                                                                label=f"Agent Type {agent_idx} - {subkey}")
215 |                                     info_axes[env_idx,subkey_to_index[subkey]].set_title(f"Environment {env_idx} - {key}.{subkey}")
216 |                                     info_axes[env_idx,subkey_to_index[subkey]].set_xlabel("Steps")
217 |                                     info_axes[env_idx,subkey_to_index[subkey]].set_ylabel(f"{key}.{subkey}")
218 |                                     info_axes[env_idx,subkey_to_index[subkey]].legend()
219 |                                     info_axes[env_idx,subkey_to_index[subkey]].grid(True)
220 |                                 else:
221 |                                     print(f"Skipping plotting for {key}.{subkey} as it has more than 2 dimensions. {subvalues.shape}")
222 |                 # Save the figure for this metric
223 |                 metric_path = os.path.join(output_dir, f"{key}_all_agents_{time.strftime('%Y%m%d-%H%M%S')}.png")
224 |                 info_fig.tight_layout()
225 |                 info_fig.savefig(metric_path)
226 |                 plt.close(info_fig)
227 |             elif hasattr(values, 'shape'):  # Handle normal array case
228 |                 # Create a figure with subplots for each environment
229 |                 info_fig, info_axes = plt.subplots(num_envs, 1, figsize=(10, 5*num_envs), sharex=True)
230 |                 if num_envs == 1:
231 |                     info_axes = [info_axes]
232 |                 # Plot this metric for each agent type on the same subplot
233 |                 for env_idx in range(num_envs):
234 |                     for agent_idx, traj in enumerate(traj_batch):
235 |                         if 'agent' in traj.info and key in traj.info['agent']:
236 |                             values = traj.info['agent'][key]      
237 |                             if len(values.shape) <= 2:  # Only plot simple scalar features
238 |                                 env_values = values[:, env_idx] if len(values.shape) > 1 else values
239 |                                 info_axes[env_idx].plot(steps, env_values, color=cmap(agent_idx), label=f"Agent Type {agent_idx}")
240 |                                 info_axes[env_idx].set_title(f"Environment {env_idx} - {key}")
241 |                                 info_axes[env_idx].set_xlabel("Steps")
242 |                                 info_axes[env_idx].set_ylabel(key)
243 |                                 info_axes[env_idx].legend()
244 |                                 info_axes[env_idx].grid(True)
245 |                             else:
246 |                                 print(f"Skipping plotting for {key} as it has more than 2 dimensions. {values.shape}")
247 |                 # Save the figure for this metric
248 |                 metric_path = os.path.join(output_dir, f"{key}_all_agents_{time.strftime('%Y%m%d-%H%M%S')}.png")
249 |                 info_fig.tight_layout()
250 |                 info_fig.savefig(metric_path)
251 |                 plt.close(info_fig)
252 |             else:
253 |                 print(f"Skipping plotting for {key} as it is neither a dict nor has a shape attribute.")
254 |     
255 |     # Plot world info
256 |     if hasattr(traj_batch[0], 'info') and 'world' in traj_batch[0].info:
257 |         world_info = traj_batch[0].info['world']
258 |         steps = np.arange(traj_batch[0].reward.shape[0])
259 |         
260 |         # Collect all unique keys from world info
261 |         world_keys = set()
262 |         for key in world_info:
263 |             world_keys.add(key)
264 |         
265 |         # For each world metric, create a figure with subplots for each environment
266 |         for key in world_keys:
267 |             if key not in features:
268 |                 continue
269 |             # Skip if the key doesn't contain plottable data
270 |             if key not in world_info or not hasattr(world_info[key], 'shape'):
271 |                 continue
272 |                 
273 |             values = world_info[key]
274 |             
275 |             # Skip complex data structures
276 |             if len(values.shape) > 3:
277 |                 print(f"Skipping plotting for world info {key} as it has more than 3 dimensions. {values.shape}")
278 |                 continue
279 |                 
280 |             # Create a figure with subplots for each environment
281 |             world_fig, world_axes = plt.subplots(num_envs, 1, figsize=(10, 5*num_envs), sharex=True)
282 |             if num_envs == 1:
283 |                 world_axes = [world_axes]
284 |             
285 |             # Plot this world metric for each environment
286 |             for env_idx in range(num_envs):
287 |                 if len(values.shape) == 2:
288 |                     env_values = values[:, env_idx]
289 |                     world_axes[env_idx].plot(steps, env_values, color='blue')  # World info uses blue
290 |                 elif len(values.shape) == 3:
291 |                     env_values = values[:, env_idx,0]+ values[:, env_idx,1]/1e9
292 |                     world_axes[env_idx].plot(steps, env_values, color='blue')  # World info uses blue
293 |                 else:
294 |                     world_axes[env_idx].plot(steps, values, color='blue')  # World info uses blue
295 |                     
296 |                 world_axes[env_idx].set_title(f"Environment {env_idx} - World {key}")
297 |                 world_axes[env_idx].set_xlabel("Steps")
298 |                 world_axes[env_idx].set_ylabel(key)
299 |                 world_axes[env_idx].grid(True)
300 |             
301 |             # Save the figure for this metric
302 |             world_metric_path = os.path.join(output_dir, f"world_{key}_{time.strftime('%Y%m%d-%H%M%S')}.png")
303 |             world_fig.tight_layout()
304 |             world_fig.savefig(world_metric_path)
305 |             plt.close(world_fig)
306 | 
307 | 
308 | 
309 | def plot_episode_features(traj_batch, output_dir="intra-episode-figs", feature_names=None, obs_features=None):
310 |     """
311 |     Plot features from trajectory batch for each environment.
312 |     
313 |     Args:
314 |         traj_batch: List of trajectory objects containing episode data
315 |         feature_names: Optional list of feature names to plot. If None, will try to infer from data.
316 |     """
317 |     # Create output directory if it doesn't exist
318 |     os.makedirs(output_dir, exist_ok=True)
319 |     
320 |     num_agent_types = len(traj_batch)
321 |     # Create a colormap that can handle all agent types
322 |     cmap = plt.cm.get_cmap('viridis', num_agent_types)
323 |     
324 |     # Get the number of environments
325 |     num_envs = traj_batch[0].reward.shape[1] if len(traj_batch[0].reward.shape) > 1 else 1
326 |     
327 |     # Create figure with subplots for each environment
328 |     fig, axes = plt.subplots(num_envs, 1, figsize=(12, 5*num_envs), sharex=True)
329 |     if num_envs == 1:
330 |         axes = [axes]
331 |     
332 |     # Plot rewards - one subplot per environment on a single figure
333 |     fig_rewards = plt.figure(figsize=(12, 5*num_envs))
334 |     for env_idx in range(num_envs):
335 |         ax_reward = fig_rewards.add_subplot(num_envs, 1, env_idx+1)
336 |         steps = np.arange(traj_batch[0].reward.shape[0])
337 |         
338 |         # For each agent type, plot a line with different color
339 |         for agent_idx, traj in enumerate(traj_batch):
340 |             rewards = traj.reward
341 |             env_rewards = rewards[:, env_idx] if len(rewards.shape) > 1 else rewards
342 |             ax_reward.plot(steps, env_rewards, color=cmap(agent_idx), label=f"Agent Type {agent_idx}")
343 |         
344 |         ax_reward.set_title(f"Environment {env_idx} Rewards")
345 |         ax_reward.set_xlabel("Steps")
346 |         ax_reward.set_ylabel("Reward")
347 |         ax_reward.legend()
348 |         ax_reward.grid(True)
349 |     
350 |     fig_rewards.tight_layout()
351 |     output_path = os.path.join(output_dir, f"rewards_all_envs_{time.strftime('%Y%m%d-%H%M%S')}.png")
352 |     fig_rewards.savefig(output_path)
353 |     plt.close(fig_rewards)
354 |     
355 |     # Plot actions - one subplot per environment on a single figure
356 |     fig_actions = plt.figure(figsize=(12, 5*num_envs))
357 |     for env_idx in range(num_envs):
358 |         ax_action = fig_actions.add_subplot(num_envs, 1, env_idx+1)
359 |         steps = np.arange(traj_batch[0].action.shape[0])
360 |         
361 |         # For each agent type, plot a line with different color
362 |         for agent_idx, traj in enumerate(traj_batch):
363 |             actions = traj.action
364 |             env_actions = actions[:, env_idx] if len(actions.shape) > 1 else actions
365 |             ax_action.plot(steps, env_actions, color=cmap(agent_idx), label=f"Agent Type {agent_idx}")
366 |         
367 |         ax_action.set_title(f"Environment {env_idx} Actions")
368 |         ax_action.set_xlabel("Steps")
369 |         ax_action.set_ylabel("Action")
370 |         ax_action.legend()
371 |         ax_action.grid(True)
372 |     
373 |     fig_actions.tight_layout()
374 |     output_path = os.path.join(output_dir, f"actions_all_envs_{time.strftime('%Y%m%d-%H%M%S')}.png")
375 |     fig_actions.savefig(output_path)
376 |     plt.close(fig_actions)
377 |     
378 |     # Plot agent-specific info
379 |     if hasattr(traj_batch[0], 'info') and traj_batch[0].info is not None:
380 |         # Collect all unique keys from agent info across all agent types
381 |         all_keys = set()
382 |         for agent_idx, traj in enumerate(traj_batch):
383 |             if 'agent' in traj.info:
384 |                 agent_info = traj.info['agent']
385 |                 all_keys.update(agent_info.keys())
386 |         print(all_keys)
387 |         
388 |         # For each metric, create a figure with subplots for each environment
389 |         for key in all_keys:
390 |             for agent_idx, traj in enumerate(traj_batch):
391 |                 if 'agent' in traj.info and key in traj.info['agent']:
392 |                     values = traj.info['agent'][key]
393 |             # If values is a dict, create a subplot for each key in the dict
394 |             if isinstance(values, dict):
395 |                 # Create a figure with subplots for each environment
396 |                 all_sub_keys = set()
397 |                 for agent_idx, traj in enumerate(traj_batch):
398 |                     if 'agent' in traj.info:
399 |                         agent_info = traj.info['agent'][key]
400 |                         all_sub_keys.update(agent_info.keys())
401 |                 print(all_sub_keys)
402 |                 subkey_to_index = {subkey: i for i, subkey in enumerate(all_sub_keys)}
403 |                 info_fig, info_axes = plt.subplots(num_envs, len(all_sub_keys), figsize=(6*len(all_sub_keys), 5*num_envs), sharex=True)
404 |                 if num_envs == 1:
405 |                     info_axes = [info_axes]
406 |                 for env_idx in range(num_envs):
407 |                     for agent_idx, traj in enumerate(traj_batch):
408 |                         if 'agent' in traj.info and key in traj.info['agent']:
409 |                             values = traj.info['agent'][key]      
410 |                             print(values.keys())
411 |                             for subkey in all_sub_keys:
412 |                                 if subkey not in values:
413 |                                     # print(f"Skipping plotting for {key}.{subkey} as it is not present in agent info.")
414 |                                     continue
415 |                                 subvalues = values[subkey]
416 |                                 if len(subvalues.shape) <= 2:  # Only plot simple scalar features
417 |                                     # print(f"Plotting {key}.{subkey} for agent type {agent_idx} in environment {env_idx}")
418 |                                     env_values = subvalues[:, env_idx] if len(subvalues.shape) > 1 else subvalues
419 |                                     info_axes[env_idx,subkey_to_index[subkey]].plot(steps, env_values, color=cmap(agent_idx), 
420 |                                                                label=f"Agent Type {agent_idx} - {subkey}")
421 |                                     info_axes[env_idx,subkey_to_index[subkey]].set_title(f"Environment {env_idx} - {key}.{subkey}")
422 |                                     info_axes[env_idx,subkey_to_index[subkey]].set_xlabel("Steps")
423 |                                     info_axes[env_idx,subkey_to_index[subkey]].set_ylabel(f"{key}.{subkey}")
424 |                                     info_axes[env_idx,subkey_to_index[subkey]].legend()
425 |                                     info_axes[env_idx,subkey_to_index[subkey]].grid(True)
426 |                                 else:
427 |                                     print(f"Skipping plotting for {key}.{subkey} as it has more than 2 dimensions. {subvalues.shape}")
428 |                 # Save the figure for this metric
429 |                 metric_path = os.path.join(output_dir, f"{key}_all_agents_{time.strftime('%Y%m%d-%H%M%S')}.png")
430 |                 info_fig.tight_layout()
431 |                 info_fig.savefig(metric_path)
432 |                 plt.close(info_fig)
433 |             elif hasattr(values, 'shape'):  # Handle normal array case
434 |                 # Create a figure with subplots for each environment
435 |                 info_fig, info_axes = plt.subplots(num_envs, 1, figsize=(10, 5*num_envs), sharex=True)
436 |                 if num_envs == 1:
437 |                     info_axes = [info_axes]
438 |                 # Plot this metric for each agent type on the same subplot
439 |                 for env_idx in range(num_envs):
440 |                     for agent_idx, traj in enumerate(traj_batch):
441 |                         if 'agent' in traj.info and key in traj.info['agent']:
442 |                             values = traj.info['agent'][key]      
443 |                             if len(values.shape) <= 2:  # Only plot simple scalar features
444 |                                 env_values = values[:, env_idx] if len(values.shape) > 1 else values
445 |                                 info_axes[env_idx].plot(steps, env_values, color=cmap(agent_idx), label=f"Agent Type {agent_idx}")
446 |                                 info_axes[env_idx].set_title(f"Environment {env_idx} - {key}")
447 |                                 info_axes[env_idx].set_xlabel("Steps")
448 |                                 info_axes[env_idx].set_ylabel(key)
449 |                                 info_axes[env_idx].legend()
450 |                                 info_axes[env_idx].grid(True)
451 |                             else:
452 |                                 print(f"Skipping plotting for {key} as it has more than 2 dimensions. {values.shape}")
453 |                 # Save the figure for this metric
454 |                 metric_path = os.path.join(output_dir, f"{key}_all_agents_{time.strftime('%Y%m%d-%H%M%S')}.png")
455 |                 info_fig.tight_layout()
456 |                 info_fig.savefig(metric_path)
457 |                 plt.close(info_fig)
458 |             else:
459 |                 print(f"Skipping plotting for {key} as it is neither a dict nor has a shape attribute.")
460 |     
461 |     # Plot world info
462 |     if hasattr(traj_batch[0], 'info') and 'world' in traj_batch[0].info:
463 |         world_info = traj_batch[0].info['world']
464 |         steps = np.arange(traj_batch[0].reward.shape[0])
465 |         
466 |         # Collect all unique keys from world info
467 |         world_keys = set()
468 |         for key in world_info:
469 |             world_keys.add(key)
470 |         
471 |         # For each world metric, create a figure with subplots for each environment
472 |         for key in world_keys:
473 |             # Skip if the key doesn't contain plottable data
474 |             if key not in world_info or not hasattr(world_info[key], 'shape'):
475 |                 continue
476 |                 
477 |             values = world_info[key]
478 |             
479 |             # Skip complex data structures
480 |             if len(values.shape) > 3:
481 |                 print(f"Skipping plotting for world info {key} as it has more than 3 dimensions. {values.shape}")
482 |                 continue
483 |                 
484 |             # Create a figure with subplots for each environment
485 |             world_fig, world_axes = plt.subplots(num_envs, 1, figsize=(10, 5*num_envs), sharex=True)
486 |             if num_envs == 1:
487 |                 world_axes = [world_axes]
488 |             
489 |             # Plot this world metric for each environment
490 |             for env_idx in range(num_envs):
491 |                 if len(values.shape) == 2:
492 |                     env_values = values[:, env_idx]
493 |                     world_axes[env_idx].plot(steps, env_values, color='blue')  # World info uses blue
494 |                 elif len(values.shape) == 3:
495 |                     env_values = values[:, env_idx,0]+ values[:, env_idx,1]/1e9
496 |                     world_axes[env_idx].plot(steps, env_values, color='blue')  # World info uses blue
497 |                 else:
498 |                     world_axes[env_idx].plot(steps, values, color='blue')  # World info uses blue
499 |                     
500 |                 world_axes[env_idx].set_title(f"Environment {env_idx} - World {key}")
501 |                 world_axes[env_idx].set_xlabel("Steps")
502 |                 world_axes[env_idx].set_ylabel(key)
503 |                 world_axes[env_idx].grid(True)
504 |             
505 |             # Save the figure for this metric
506 |             world_metric_path = os.path.join(output_dir, f"world_{key}_{time.strftime('%Y%m%d-%H%M%S')}.png")
507 |             world_fig.tight_layout()
508 |             world_fig.savefig(world_metric_path)
509 |             plt.close(world_fig)
510 | 
511 | 
512 | if __name__ == "__main__":
513 |     main()
514 | 


--------------------------------------------------------------------------------
/gymnax_exchange/jaxrl/MARL/baseline_eval/plotting_combinations.py:
--------------------------------------------------------------------------------
  1 | #!/usr/bin/env python3
  2 | """
  3 | Plotting script for visualizing agent rewards in Baseline vs Learned configurations.
  4 | Creates an n×n heatmap where n is the number of agents, showing rewards for each agent
  5 | across all possible combinations of Baseline (B) and Learned (L) agent types.
  6 | 
  7 | For example, with 2 agents:
  8 | - Config (0,0): Baseline vs Baseline
  9 | - Config (0,1): Baseline vs Learned  
 10 | - Config (1,0): Learned vs Baseline
 11 | - Config (1,1): Learned vs Learned
 12 | 
 13 | Higher values (green) indicate better performance, lower values (red) indicate worse performance.
 14 | """
 15 | 
 16 | import argparse
 17 | import numpy as np
 18 | import matplotlib.pyplot as plt
 19 | import matplotlib.figure as mpl_figure
 20 | import matplotlib.colors as mcolors
 21 | import matplotlib.patches as mpatches
 22 | import matplotlib.lines as mlines
 23 | import seaborn as sns
 24 | import pickle
 25 | import json
 26 | import os
 27 | from pathlib import Path
 28 | from typing import Dict, List, Tuple, Optional, Union
 29 | import pandas as pd
 30 | 
 31 | 
 32 | class AgentRewardPlotter:
 33 |     """Class for creating heatmap visualizations of agent rewards in Baseline vs Learned configurations."""
 34 |     
 35 |     def __init__(self, n_agents: int = 2, figsize: Tuple[int, int] = (12, 10)):
 36 |         """
 37 |         Initialize the plotter.
 38 |         
 39 |         Args:
 40 |             n_agents: Number of agents (creates 2^n_agents configurations)
 41 |             figsize: Figure size for the plot
 42 |         """
 43 |         self.n_agents = n_agents
 44 |         self.n_configs = 2 ** n_agents  # 2^n possible configurations
 45 |         self.figsize = figsize
 46 |         self.data = None
 47 |         
 48 |     def generate_sample_data(self) -> Dict[str, np.ndarray]:
 49 |         """
 50 |         Generate sample reward data for testing.
 51 |         
 52 |         Returns:
 53 |             Dictionary containing reward data for each configuration
 54 |         """
 55 |         np.random.seed(42)  # For reproducibility
 56 |         
 57 |         data = {}
 58 |         
 59 |         # Generate all possible configurations (Baseline=0, Learned=1)
 60 |         for config_id in range(2 ** self.n_agents):
 61 |             # Convert config_id to binary representation for agent types
 62 |             agent_types = []
 63 |             temp_id = config_id
 64 |             for _ in range(self.n_agents):
 65 |                 agent_types.append(temp_id % 2)  # 0=Baseline, 1=Learned
 66 |                 temp_id //= 2
 67 |             
 68 |             # Create configuration name
 69 |             type_names = ['B' if t == 0 else 'L' for t in agent_types]
 70 |             config_name = f"Config_{''.join(type_names)}"
 71 |             
 72 |             # Generate sample reward data
 73 |             # Shape: (num_episodes, n_agents)
 74 |             num_episodes = np.random.randint(50, 200)
 75 |             base_reward = np.random.uniform(0.1, 0.5, size=(num_episodes, self.n_agents))
 76 |             
 77 |             # Different reward patterns for Baseline vs Learned agents
 78 |             for agent_idx in range(self.n_agents):
 79 |                 if agent_types[agent_idx] == 0:  # Baseline agent
 80 |                     # Baseline: steady but lower performance
 81 |                     improvement = np.linspace(1.0, 1.5, num_episodes)
 82 |                     base_reward[:, agent_idx] = base_reward[:, agent_idx] * improvement
 83 |                 else:  # Learned agent
 84 |                     # Learned: higher performance with learning curve
 85 |                     improvement = np.linspace(1.0, 3.0, num_episodes)
 86 |                     base_reward[:, agent_idx] = base_reward[:, agent_idx] * improvement
 87 |                     
 88 |                     # Add some volatility during learning
 89 |                     volatility = 1 + 0.2 * np.sin(np.linspace(0, 8, num_episodes))
 90 |                     base_reward[:, agent_idx] = base_reward[:, agent_idx] * volatility
 91 |             
 92 |             # Add some noise and ensure positive rewards
 93 |             reward_data = base_reward + np.random.normal(0, 0.05, base_reward.shape)
 94 |             reward_data = np.maximum(reward_data, 0.01)  # Ensure positive rewards
 95 |             
 96 |             data[config_name] = reward_data
 97 |         
 98 |         self.data = data
 99 |         return data
100 |     
101 |     def load_data(self, data_path: str) -> Dict[str, np.ndarray]:
102 |         """
103 |         Load reward data from file.
104 |         
105 |         Args:
106 |             data_path: Path to the data file (.pkl, .json, .npz, or .csv)
107 |             
108 |         Returns:
109 |             Dictionary containing reward data for each configuration
110 |         """
111 |         data_path_obj = Path(data_path)
112 |         
113 |         if not data_path_obj.exists():
114 |             raise FileNotFoundError(f"Data file not found: {data_path_obj}")
115 |         
116 |         if data_path_obj.suffix == '.pkl':
117 |             with open(data_path_obj, 'rb') as f:
118 |                 data = pickle.load(f)
119 |         elif data_path_obj.suffix == '.json':
120 |             with open(data_path_obj, 'r') as f:
121 |                 json_data = json.load(f)
122 |             # Convert lists back to numpy arrays
123 |             data = {k: np.array(v) for k, v in json_data.items()}
124 |         elif data_path_obj.suffix == '.npz':
125 |             npz_data = np.load(data_path_obj)
126 |             data = {key: npz_data[key] for key in npz_data.files}
127 |         elif data_path_obj.suffix == '.csv':
128 |             # Load CSV data
129 |             df = pd.read_csv(data_path_obj)
130 |             data = self._parse_csv_data(df)
131 |         else:
132 |             raise ValueError(f"Unsupported file format: {data_path_obj.suffix}")
133 |         
134 |         self.data = data
135 |         return data
136 |     
137 |     def _parse_csv_data(self, df: pd.DataFrame) -> Dict[str, np.ndarray]:
138 |         """
139 |         Parse CSV data into the expected format.
140 |         
141 |         Expected CSV format:
142 |         - Column 'Configuration': Configuration names (e.g., Config_BB, Config_BL, Config_LB, Config_LL)
143 |         - Column 'Episode': Episode numbers
144 |         - Columns 'MarketMaker', 'ExecutionAgent', etc.: Reward values for each agent
145 |         
146 |         Args:
147 |             df: DataFrame loaded from CSV
148 |             
149 |         Returns:
150 |             Dictionary containing reward data for each configuration
151 |         """
152 |         data = {}
153 |         
154 |         # Group by configuration
155 |         for config_name, group in df.groupby('Configuration'):
156 |             # Sort by episode to ensure correct order
157 |             group = group.sort_values('Episode')
158 |             
159 |             # Extract agent columns - try both naming conventions
160 |             agent_cols = []
161 |             
162 |             # First try MarketMaker, ExecutionAgent naming
163 |             if 'MarketMaker' in group.columns:
164 |                 agent_cols.append('MarketMaker')
165 |             if 'ExecutionAgent' in group.columns:
166 |                 agent_cols.append('ExecutionAgent')
167 |             
168 |             # If we don't have the named columns, fall back to Agent_0, Agent_1, etc.
169 |             if not agent_cols:
170 |                 agent_cols = [col for col in group.columns if col.startswith('Agent_')]
171 |             
172 |             if not agent_cols:
173 |                 raise ValueError("No agent columns found. Expected columns like 'MarketMaker', 'ExecutionAgent' or 'Agent_0', 'Agent_1', etc.")
174 |             
175 |             # Convert to numpy array (episodes x agents)
176 |             reward_data = group[agent_cols].values
177 |             data[config_name] = reward_data
178 |         
179 |         return data
180 |     
181 |     def _create_sample_csv_data(self) -> pd.DataFrame:
182 |         """
183 |         Create sample data in CSV format for easy editing.
184 |         
185 |         Returns:
186 |             DataFrame with sample reward data
187 |         """
188 |         # Ensure we have data to work with
189 |         if self.data is None:
190 |             # Generate sample data if none exists
191 |             self.generate_sample_data()
192 |         
193 |         if self.data is None:
194 |             raise ValueError("Failed to generate sample data")
195 |         
196 |         # Define agent column names based on number of agents
197 |         if self.n_agents == 2:
198 |             agent_names = ['MarketMaker', 'ExecutionAgent']
199 |         else:
200 |             # For more than 2 agents, use Agent_0, Agent_1, etc.
201 |             agent_names = [f'Agent_{i}' for i in range(self.n_agents)]
202 |         
203 |         rows = []
204 |         for config_name, reward_data in self.data.items():
205 |             num_episodes, num_agents = reward_data.shape
206 |             for episode in range(num_episodes):
207 |                 row = {
208 |                     'Configuration': config_name,
209 |                     'Episode': episode
210 |                 }
211 |                 for agent_idx in range(num_agents):
212 |                     if agent_idx < len(agent_names):
213 |                         row[agent_names[agent_idx]] = reward_data[episode, agent_idx]
214 |                     else:
215 |                         row[f'Agent_{agent_idx}'] = reward_data[episode, agent_idx]
216 |                 rows.append(row)
217 |         
218 |         return pd.DataFrame(rows)
219 |     
220 |     def save_data(self, data: Dict[str, np.ndarray], save_path: str):
221 |         """
222 |         Save reward data to file.
223 |         
224 |         Args:
225 |             data: Dictionary containing reward data
226 |             save_path: Path where to save the data
227 |         """
228 |         save_path_obj = Path(save_path)
229 |         save_path_obj.parent.mkdir(parents=True, exist_ok=True)
230 |         
231 |         if save_path_obj.suffix == '.pkl':
232 |             with open(save_path_obj, 'wb') as f:
233 |                 pickle.dump(data, f)
234 |         elif save_path_obj.suffix == '.json':
235 |             # Convert numpy arrays to lists for JSON serialization
236 |             json_data = {k: v.tolist() for k, v in data.items()}
237 |             with open(save_path_obj, 'w') as f:
238 |                 json.dump(json_data, f, indent=2)
239 |         elif save_path_obj.suffix == '.npz':
240 |             np.savez(save_path_obj, **data)
241 |         elif save_path_obj.suffix == '.csv':
242 |             # Save as CSV for easy editing
243 |             df = self._create_sample_csv_data()
244 |             df.to_csv(save_path_obj, index=False)
245 |         else:
246 |             raise ValueError(f"Unsupported file format: {save_path_obj.suffix}")
247 |     
248 |     def compute_overall_reward(self, reward_data: np.ndarray, method: str = 'mean') -> Union[float, np.floating]:
249 |         """
250 |         Compute overall reward from agent reward data.
251 |         
252 |         Args:
253 |             reward_data: Array of shape (episodes, agents)
254 |             method: Method to compute overall reward ('mean', 'sum', 'max', 'final')
255 |             
256 |         Returns:
257 |             Overall reward value
258 |         """
259 |         if method == 'mean':
260 |             return np.mean(reward_data)
261 |         elif method == 'sum':
262 |             return np.sum(reward_data, axis=1).mean()  # Sum across agents, then average across episodes
263 |         elif method == 'max':
264 |             return np.max(reward_data)
265 |         elif method == 'final':
266 |             return np.mean(reward_data[-1, :])  # Final episode average
267 |         else:
268 |             raise ValueError(f"Unknown method: {method}")
269 |     
270 |     def create_heatmap(self, method: str = 'mean', cmap: str = 'RdYlGn', 
271 |                       title: Optional[str] = None, save_path: Optional[str] = None) -> mpl_figure.Figure:
272 |         """
273 |         Create heatmap visualization with triangular splits showing each agent's rewards.
274 |         
275 |         Args:
276 |             method: Method to compute overall reward
277 |             cmap: Colormap for the heatmap (default: RdYlGn for red-green)
278 |             title: Custom title for the plot
279 |             save_path: Path to save the plot
280 |             
281 |         Returns:
282 |             Matplotlib figure object
283 |         """
284 |         if self.data is None:
285 |             raise ValueError("No data loaded. Call generate_sample_data() or load_data() first.")
286 |         
287 |         # Create grid size based on number of agents (2^n_agents configurations arranged in square)
288 |         grid_size = int(np.ceil(np.sqrt(self.n_configs)))
289 |         
290 |         # Separate matrices for each agent
291 |         agent_rewards = []
292 |         for agent_idx in range(self.n_agents):
293 |             agent_rewards.append(np.full((grid_size, grid_size), np.nan))
294 |         
295 |         config_labels = []
296 |         
297 |         # Fill matrices with configuration data
298 |         # For 2x2 grid, we want: BB bottom-left, BL bottom-right, LB top-left, LL top-right
299 |         config_positions = {
300 |             'Config_BB': (1, 0),  # bottom-left
301 |             'Config_BL': (1, 1),  # bottom-right  
302 |             'Config_LB': (0, 0),  # top-left
303 |             'Config_LL': (0, 1)   # top-right
304 |         }
305 |         
306 |         for config_name, reward_data in self.data.items():
307 |             if config_name in config_positions:
308 |                 row, col = config_positions[config_name]
309 |                 config_labels.append(config_name)
310 |                 
311 |                 print(f"Debug: {config_name} -> grid position ({row}, {col})")
312 |                 print(f"Debug: {config_name} shape: {reward_data.shape}, sample data: {reward_data[:2, :]}")
313 |                 
314 |                 # Compute rewards for each agent separately
315 |                 for agent_idx in range(min(self.n_agents, reward_data.shape[1])):
316 |                     if method == 'mean':
317 |                         agent_reward = np.mean(reward_data[:, agent_idx])
318 |                     elif method == 'sum':
319 |                         agent_reward = np.sum(reward_data[:, agent_idx])
320 |                     elif method == 'max':
321 |                         agent_reward = np.max(reward_data[:, agent_idx])
322 |                     elif method == 'final':
323 |                         agent_reward = reward_data[-1, agent_idx]
324 |                     else:
325 |                         agent_reward = np.mean(reward_data[:, agent_idx])
326 |                     
327 |                     print(f"Debug: {config_name} Agent {agent_idx} reward: {agent_reward}")
328 |                     agent_rewards[agent_idx][row, col] = agent_reward
329 |             else:
330 |                 print(f"Warning: Configuration {config_name} not found in position mapping")
331 |         
332 |         # Create the figure with triangular split visualization
333 |         fig, ax = plt.subplots(figsize=self.figsize)
334 |         
335 |         # Agent names for legend
336 |         if self.n_agents == 2:
337 |             agent_names = ['MarketMaker', 'ExecutionAgent']
338 |         else:
339 |             agent_names = [f'Agent_{i}' for i in range(self.n_agents)]
340 |         
341 |         # Create triangular heatmap for 2 agents
342 |         if self.n_agents == 2:
343 |             # Normalize the data for consistent color scaling
344 |             all_values = []
345 |             for agent_idx in range(2):
346 |                 valid_values = agent_rewards[agent_idx][~np.isnan(agent_rewards[agent_idx])]
347 |                 all_values.extend(valid_values)
348 |             
349 |             # Fixed color scale from -50 to 0
350 |             vmin, vmax = -25, 0
351 |             
352 |             # Create triangular plots using the fixed position mapping
353 |             config_positions = {
354 |                 'Config_BB': (1, 0),  # bottom-left
355 |                 'Config_BL': (1, 1),  # bottom-right  
356 |                 'Config_LB': (0, 0),  # top-left
357 |                 'Config_LL': (0, 1)   # top-right
358 |             }
359 |             
360 |             for config_name in config_positions:
361 |                 if config_name in self.data:
362 |                     row, col = config_positions[config_name]
363 |                     
364 |                     # Get values for both agents
365 |                     mm_val = agent_rewards[0][row, col]  # MarketMaker
366 |                     ea_val = agent_rewards[1][row, col]  # ExecutionAgent
367 |                     
368 |                     if not (np.isnan(mm_val) or np.isnan(ea_val)):
369 |                         # Create triangular patches
370 |                         # Upper triangle for MarketMaker
371 |                         triangle_upper = mpatches.Polygon([(col, row+1), (col+1, row+1), (col+1, row)], 
372 |                                                    closed=True)
373 |                         # Lower triangle for ExecutionAgent  
374 |                         triangle_lower = mpatches.Polygon([(col, row), (col, row+1), (col+1, row)], 
375 |                                                    closed=True)
376 |                             
377 |                     # Fixed color scale from -50 to 0 with saturation for out-of-range values
378 |                     color_min, color_max = -50, 0
379 |                     
380 |                     # Clip values to color range for color mapping
381 |                     mm_val_clipped = np.clip(mm_val, color_min, color_max)
382 |                     ea_val_clipped = np.clip(ea_val, color_min, color_max)
383 |                     
384 |                     # Normalize colors
385 |                     colormap = plt.colormaps.get_cmap(cmap)
386 |                     mm_color = colormap((mm_val_clipped - color_min) / (color_max - color_min))
387 |                     ea_color = colormap((ea_val_clipped - color_min) / (color_max - color_min))
388 |                     
389 |                     triangle_upper.set_facecolor(mm_color)
390 |                     triangle_lower.set_facecolor(ea_color)
391 |                     triangle_upper.set_edgecolor('white')
392 |                     triangle_lower.set_edgecolor('white')
393 |                     triangle_upper.set_linewidth(1)
394 |                     triangle_lower.set_linewidth(1)
395 |                     
396 |                     ax.add_patch(triangle_upper)
397 |                     ax.add_patch(triangle_lower)
398 |                     
399 |                     # Add text annotations with original values
400 |                     ax.text(col+0.75, row+0.75, f'{mm_val:.1f}', ha='center', va='center', 
401 |                            fontsize=8, fontweight='bold')
402 |                     ax.text(col+0.25, row+0.25, f'{ea_val:.1f}', ha='center', va='center', 
403 |                            fontsize=8, fontweight='bold')
404 |             
405 |             # Set up the plot
406 |             ax.set_xlim(0, grid_size)
407 |             ax.set_ylim(0, grid_size)
408 |             ax.set_aspect('equal')
409 |             # Orientation is now handled by flipped grid mapping
410 |             
411 |             # Create custom labels for individual agents
412 |             if self.n_agents == 2:
413 |                 # For 2 agents, create a 2x2 grid where:
414 |                 # X-axis represents ExecutionAgent (Agent 1): Baseline/Learned
415 |                 # Y-axis represents MarketMaker (Agent 0): Baseline/Learned
416 |                 config_labels_x = ['Baseline', 'Learned']  # ExecutionAgent
417 |                 config_labels_y = ['Baseline', 'Learned']  # MarketMaker
418 |             else:
419 |                 # For more agents, use the combined configuration labels
420 |                 config_labels_x = []
421 |                 config_labels_y = []
422 |                 for i in range(grid_size):
423 |                     for j in range(grid_size):
424 |                         config_id = i * grid_size + j
425 |                         if config_id < self.n_configs:
426 |                             agent_types = []
427 |                             temp_id = config_id
428 |                             for _ in range(self.n_agents):
429 |                                 agent_types.append(temp_id % 2)
430 |                                 temp_id //= 2
431 |                             type_names = ['B' if t == 0 else 'L' for t in agent_types]
432 |                             label = ''.join(type_names)
433 |                         else:
434 |                             label = ''
435 |                         
436 |                         if i == 0:
437 |                             config_labels_x.append(label)
438 |                         if j == 0:
439 |                             config_labels_y.append(label)
440 |             
441 |             # Set ticks and labels
442 |             ax.set_xticks([i + 0.5 for i in range(len(config_labels_x))])
443 |             ax.set_xticklabels(config_labels_x)
444 |             ax.set_yticks([i + 0.5 for i in range(len(config_labels_y))])
445 |             ax.set_yticklabels(config_labels_y)
446 |             
447 |             # Add colorbar
448 |             sm = plt.cm.ScalarMappable(cmap=cmap, norm=mcolors.Normalize(vmin=vmin, vmax=vmax))
449 |             sm.set_array([])
450 |             cbar = plt.colorbar(sm, ax=ax)
451 |             cbar.set_label(f'Reward ({method}) [Range: -50 to 0, saturated beyond]', rotation=270, labelpad=15)
452 |             
453 |             # Add legend for triangles
454 |             legend_elements = [
455 |                 mlines.Line2D([0], [0], marker='^', color='w', markerfacecolor='gray', 
456 |                           markersize=10, label=f'{agent_names[0]} (upper triangle)'),
457 |                 mlines.Line2D([0], [0], marker='v', color='w', markerfacecolor='gray', 
458 |                           markersize=10, label=f'{agent_names[1]} (lower triangle)')
459 |             ]
460 |             ax.legend(handles=legend_elements, loc='upper left', bbox_to_anchor=(1.15, 1))
461 |             
462 |         else:
463 |             # For more than 2 agents, fall back to regular heatmap with average
464 |             combined_rewards = np.full((grid_size, grid_size), np.nan)
465 |             for i in range(grid_size):
466 |                 for j in range(grid_size):
467 |                     valid_rewards = []
468 |                     for agent_idx in range(self.n_agents):
469 |                         if not np.isnan(agent_rewards[agent_idx][i, j]):
470 |                             valid_rewards.append(agent_rewards[agent_idx][i, j])
471 |                     if valid_rewards:
472 |                         combined_rewards[i, j] = np.mean(valid_rewards)
473 |             
474 |             mask = np.isnan(combined_rewards)
475 |             sns.heatmap(combined_rewards, annot=~mask, fmt='.3f', cmap=cmap, ax=ax,
476 |                        mask=mask, cbar_kws={'label': f'Average Reward ({method})'})
477 |         
478 |         # Set title
479 |         if title is None:
480 |             if self.n_agents == 2:
481 |                 title = f'Agent Rewards: Baseline vs Learned\n(Upper: {agent_names[0]}, Lower: {agent_names[1]})'
482 |             else:
483 |                 title = f'Agent Rewards: Baseline vs Learned ({self.n_agents} Agents, {self.n_configs} Configurations)'
484 |         ax.set_title(title, fontsize=14, pad=20)
485 |         
486 |         # Labels
487 |         if self.n_agents == 2:
488 |             ax.set_xlabel('ExecutionAgent Type', fontsize=12)
489 |             ax.set_ylabel('MarketMaker Type', fontsize=12)
490 |         else:
491 |             ax.set_xlabel('Agent Configuration (B=Baseline, L=Learned)', fontsize=12)
492 |             ax.set_ylabel('Agent Configuration (B=Baseline, L=Learned)', fontsize=12)
493 |         
494 |         plt.tight_layout()
495 |         
496 |         # Save if path provided
497 |         if save_path:
498 |             save_path_obj = Path(save_path)
499 |             save_path_obj.parent.mkdir(parents=True, exist_ok=True)
500 |             plt.savefig(save_path_obj, dpi=300, bbox_inches='tight')
501 |             print(f"Plot saved to: {save_path_obj}")
502 |         
503 |         return fig
504 |     
505 |     def create_detailed_plot(self, save_path: Optional[str] = None) -> mpl_figure.Figure:
506 |         """
507 |         Create a detailed plot showing individual agent rewards and overall reward heatmap.
508 |         
509 |         Args:
510 |             save_path: Path to save the plot
511 |             
512 |         Returns:
513 |             Matplotlib figure object
514 |         """
515 |         if self.data is None:
516 |             raise ValueError("No data loaded. Call generate_sample_data() or load_data() first.")
517 |         
518 |         # Create subplot layout
519 |         fig = plt.figure(figsize=(16, 12))
520 |         
521 |         # Main heatmap (top half)
522 |         ax_main = plt.subplot2grid((3, 2), (0, 0), colspan=2, rowspan=2)
523 |         
524 |         # Create grid size based on number of agents
525 |         grid_size = int(np.ceil(np.sqrt(self.n_configs)))
526 |         reward_matrix = np.full((grid_size, grid_size), np.nan)
527 |         
528 |         for config_id in range(self.n_configs):
529 |             # Use flipped row mapping for consistency
530 |             row = (grid_size - 1) - (config_id // grid_size)  # Flip row0
531 |             col = config_id % grid_size
532 |             
533 |             # Convert config_id to agent types
534 |             agent_types = []
535 |             temp_id = config_id
536 |             for _ in range(self.n_agents):
537 |                 agent_types.append(temp_id % 2)
538 |                 temp_id //= 2
539 |             
540 |             type_names = ['B' if t == 0 else 'L' for t in agent_types]
541 |             config_name = f"Config_{''.join(type_names)}"
542 |             
543 |             if config_name in self.data:
544 |                 reward_data = self.data[config_name]
545 |                 reward_matrix[row, col] = self.compute_overall_reward(reward_data, 'mean')
546 |         
547 |         # Create labels for the heatmap
548 |         row_labels = []
549 |         col_labels = []
550 |         for i in range(grid_size):
551 |             for j in range(grid_size):
552 |                 config_id = i * grid_size + j
553 |                 if config_id < self.n_configs:
554 |                     agent_types = []
555 |                     temp_id = config_id
556 |                     for _ in range(self.n_agents):
557 |                         agent_types.append(temp_id % 2)
558 |                         temp_id //= 2
559 |                     type_names = ['B' if t == 0 else 'L' for t in agent_types]
560 |                     label = ''.join(type_names)
561 |                 else:
562 |                     label = ''
563 |                 
564 |                 if j == 0:
565 |                     row_labels.append(label)
566 |                 if i == 0:
567 |                     col_labels.append(label)
568 |         
569 |         # Create main heatmap
570 |         mask = np.isnan(reward_matrix)
571 |         sns.heatmap(reward_matrix, annot=~mask, fmt='.3f', cmap='RdYlGn', ax=ax_main,
572 |                    mask=mask, cbar_kws={'label': 'Overall Reward (mean)'},
573 |                    xticklabels=col_labels, yticklabels=row_labels)
574 |         ax_main.set_title(f'Baseline vs Learned Agent Rewards ({self.n_agents} Agents, {self.n_configs} Configurations)', 
575 |                          fontsize=16, pad=20)
576 |         
577 |         # Example individual agent rewards (bottom left)
578 |         ax_agents = plt.subplot2grid((3, 2), (2, 0))
579 |         # Show first configuration as example
580 |         first_config = list(self.data.keys())[0]
581 |         if first_config in self.data:
582 |             reward_data = self.data[first_config]
583 |             # Define agent names
584 |             if self.n_agents == 2:
585 |                 agent_names = ['MarketMaker', 'ExecutionAgent']
586 |             else:
587 |                 agent_names = [f'Agent {i}' for i in range(self.n_agents)]
588 |             
589 |             for agent_idx in range(self.n_agents):
590 |                 agent_name = agent_names[agent_idx] if agent_idx < len(agent_names) else f'Agent {agent_idx}'
591 |                 ax_agents.plot(reward_data[:, agent_idx], label=agent_name, alpha=0.7)
592 |             ax_agents.set_xlabel('Episode')
593 |             ax_agents.set_ylabel('Reward')
594 |             ax_agents.set_title(f'Individual Agent Rewards - {first_config}')
595 |             ax_agents.legend()
596 |             ax_agents.grid(True, alpha=0.3)
597 |         
598 |         # Reward distribution (bottom right)
599 |         ax_dist = plt.subplot2grid((3, 2), (2, 1))
600 |         all_rewards = []
601 |         for config_data in self.data.values():
602 |             all_rewards.extend(config_data.flatten())
603 |         ax_dist.hist(all_rewards, bins=30, alpha=0.7, edgecolor='black')
604 |         ax_dist.set_xlabel('Reward Value')
605 |         ax_dist.set_ylabel('Frequency')
606 |         ax_dist.set_title('Reward Distribution Across All Configurations')
607 |         ax_dist.grid(True, alpha=0.3)
608 |         
609 |         plt.tight_layout()
610 |         
611 |         # Save if path provided
612 |         if save_path:
613 |             save_path_obj = Path(save_path)
614 |             save_path_obj.parent.mkdir(parents=True, exist_ok=True)
615 |             plt.savefig(save_path_obj, dpi=300, bbox_inches='tight')
616 |             print(f"Detailed plot saved to: {save_path_obj}")
617 |         
618 |         return fig
619 | 
620 | 
621 | def main():
622 |     """Main function with CLI interface."""
623 |     parser = argparse.ArgumentParser(description='Plot agent rewards for Baseline vs Learned configurations')
624 |     
625 |     # Data arguments
626 |     parser.add_argument('--data-path', type=str, help='Path to load data from (.pkl, .json, .npz, or .csv)')
627 |     parser.add_argument('--save-data', type=str, help='Path to save generated sample data')
628 |     parser.add_argument('--output-dir', type=str, default='./plots', 
629 |                        help='Directory to save plots (default: ./plots)')
630 |     
631 |     # Configuration arguments
632 |     parser.add_argument('--n-agents', type=int, default=2, help='Number of agents (creates 2^n configurations)')
633 |     parser.add_argument('--method', type=str, default='mean', 
634 |                        choices=['mean', 'sum', 'max', 'final'],
635 |                        help='Method to compute overall reward')
636 |     
637 |     # Plot arguments
638 |     parser.add_argument('--cmap', type=str, default='RdYlGn', help='Colormap for heatmap (default: RdYlGn)')
639 |     parser.add_argument('--title', type=str, help='Custom title for the plot')
640 |     parser.add_argument('--figsize', type=int, nargs=2, default=[12, 10], 
641 |                        help='Figure size (width height)')
642 |     parser.add_argument('--detailed', action='store_true', 
643 |                        help='Create detailed plot with individual agent rewards')
644 |     
645 |     # Parse arguments
646 |     args = parser.parse_args()
647 |     
648 |     # Create plotter
649 |     plotter = AgentRewardPlotter(n_agents=args.n_agents, figsize=tuple(args.figsize))
650 |     
651 |     # Load or generate data
652 |     if args.data_path:
653 |         print(f"Loading data from: {args.data_path}")
654 |         data = plotter.load_data(args.data_path)
655 |         print(f"Loaded data for {len(data)} configurations")
656 |     else:
657 |         print("Generating sample data...")
658 |         data = plotter.generate_sample_data()
659 |         print(f"Generated sample data for {len(data)} configurations")
660 |     
661 |     # Save data if requested
662 |     if args.save_data:
663 |         plotter.save_data(data, args.save_data)
664 |         print(f"Data saved to: {args.save_data}")
665 |     
666 |     # Create output directory
667 |     output_dir = Path(args.output_dir)
668 |     output_dir.mkdir(parents=True, exist_ok=True)
669 |     
670 |     # Create and save plots
671 |     if args.detailed:
672 |         # Create detailed plot
673 |         print("Creating detailed plot...")
674 |         fig = plotter.create_detailed_plot(
675 |             save_path=str(output_dir / f'detailed_baseline_vs_learned_n{args.n_agents}.png')
676 |         )
677 |     else:
678 |         # Create simple heatmap
679 |         print("Creating heatmap...")
680 |         fig = plotter.create_heatmap(
681 |             method=args.method,
682 |             cmap=args.cmap,
683 |             title=args.title,
684 |             save_path=str(output_dir / f'baseline_vs_learned_heatmap_n{args.n_agents}_{args.method}.png')
685 |         )
686 |     
687 |     # Show plot
688 |     plt.show()
689 |     
690 |     print("Done!")
691 | 
692 | 
693 | if __name__ == "__main__":
694 |     main()
695 | 


--------------------------------------------------------------------------------