kozax 0.0.1__tar.gz

kozax-0.0.1/.gitignore

@@ -0,0 +1 @@
+*.pyc

kozax-0.0.1/Kozax/environments/SR_environments/lorenz_attractor.py

@@ -0,0 +1,82 @@
+"""
+Kozax: Genetic programming framework in JAX
+
+Copyright (c) 2024 sdevries0
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with this program.  If not, see <https://www.gnu.org/licenses/>.
+"""
+
+import jax
+import jax.numpy as jnp
+import jax.random as jrandom
+from Kozax.environments.SR_environments.time_series_environment_base import EnvironmentBase
+
+class LorenzAttractor(EnvironmentBase):
+    def __init__(self, process_noise, obs_noise, n_obs=3):
+        n_var = 3
+        super().__init__(process_noise, obs_noise, n_var, n_obs)
+
+        self.init_mu = jnp.array([1,1,1])
+        self.init_sd = 1
+
+        self.sigma = 10
+        self.rho = 28
+        self.beta = 8/3
+        self.V = self.process_noise * jnp.eye(self.n_var)
+        self.W = self.obs_noise * jnp.eye(self.n_obs)[:self.n_obs]
+        self.C = jnp.eye(self.n_var)[:self.n_obs]
+
+    def sample_init_states(self, batch_size, key):
+        return self.init_mu + self.init_sd*jrandom.normal(key, shape=(batch_size,3))
+        # return jnp.ones((batch_size, 3))
+    
+    def drift(self, t, state, args):
+        return jnp.array([self.sigma*(state[1]-state[0]), state[0]*(self.rho-state[2])-state[1], state[0]*state[1]-self.beta*state[2]])
+
+    def diffusion(self, t, state, args):
+        return self.V
+
+    def terminate_event(self, state, **kwargs):
+        return False
+    
+class Lorenz96(EnvironmentBase):
+    def __init__(self, process_noise, obs_noise, n_dim = 3, n_obs=3):
+        n_var = n_dim
+        super().__init__(process_noise, obs_noise, n_var, n_obs)
+
+        self.F = 8
+        self.init_mu = jnp.ones(self.n_var)*self.F*0
+        self.init_sd = 1
+        
+
+        self.V = self.process_noise * jnp.eye(self.n_var)
+        self.W = self.obs_noise * jnp.eye(self.n_obs)[:self.n_obs]
+        self.C = jnp.eye(self.n_var)[:self.n_obs]
+
+    def sample_init_states(self, batch_size, key):
+        return self.init_mu + self.init_sd*jrandom.normal(key, shape=(batch_size,self.n_var))
+        # return jnp.ones((batch_size, 3))
+    
+    def sample_init_state2(self, ys, batch_size, key):
+        return ys[jrandom.choice(key, jnp.arange(ys.shape[0]), shape=(batch_size,), replace=False)]
+    
+    def drift(self, t, state, args):
+        f = lambda x_cur, x_next, x_prev1, x_prev2: (x_next - x_prev2) * x_prev1 - x_cur + self.F
+        return jax.vmap(f)(state, jnp.roll(state, -1), jnp.roll(state, 1), jnp.roll(state, 2))
+
+    def diffusion(self, t, state, args):
+        return self.V
+
+    def terminate_event(self, state, **kwargs):
+        return False

kozax-0.0.1/Kozax/environments/SR_environments/lotka_volterra.py

@@ -0,0 +1,103 @@
+"""
+Kozax: Genetic programming framework in JAX
+
+Copyright (c) 2024 sdevries0
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with this program.  If not, see <https://www.gnu.org/licenses/>.
+"""
+
+import jax
+import jax.numpy as jnp
+import jax.random as jrandom
+from Kozax.environments.SR_environments.time_series_environment_base import EnvironmentBase
+
+class LotkaVolterra(EnvironmentBase):
+    def __init__(self, process_noise, obs_noise, n_obs=2):
+        n_var = 2
+        super().__init__(process_noise, obs_noise, n_var, n_obs)
+
+        self.init_mu = jnp.array([10, 10])
+        self.init_sd = 2
+
+        self.alpha = 1.1
+        self.beta = 0.4
+        self.delta = 0.1
+        self.gamma = 0.4
+        self.V = self.process_noise * jnp.eye(self.n_var)
+        self.W = self.obs_noise * jnp.eye(self.n_obs)[:self.n_obs]
+        self.C = jnp.eye(self.n_var)[:self.n_obs]
+
+    def sample_init_states(self, batch_size, key):
+        return jrandom.uniform(key, shape = (batch_size,2), minval=5, maxval=15)
+    
+    def sample_init_state2(self, ys, batch_size, key):
+        return ys[jrandom.choice(key, jnp.arange(ys.shape[0]), shape=(batch_size,), replace=False)]
+    
+    def drift(self, t, state, args):
+        return jnp.array([self.alpha * state[0] - self.beta * state[0] * state[1], self.delta * state[0] * state[1] - self.gamma * state[1]])
+
+    def diffusion(self, t, state, args):
+        return self.V
+
+    def terminate_event(self, state, **kwargs):
+        return False
+    
+
+class LotkaVolterraN(EnvironmentBase):
+    def __init__(self, key, process_noise, obs_noise, n_dim=1, n_obs=2):
+        n_var = 2*n_dim
+        n_obs = n_dim * n_obs
+        super().__init__(process_noise, obs_noise, n_var, n_obs)
+
+        self.init_mu = 10 * jnp.ones(n_var)
+        self.init_sd = 2
+
+        keys = jrandom.split(key, (2+n_var,))
+
+        # self.rate = jrandom.uniform(keys[0], shape=n_var, minval=0.2, maxval=0.8)*jnp.repeat(jnp.array([[1.0,-1.0]]), n_dim,axis=0).ravel()
+        prey_rates = jrandom.uniform(keys[0], shape=n_dim, minval=0.5, maxval=1.0)
+        predator_rates = jrandom.uniform(keys[1], shape=n_dim, minval=-0.5, maxval=-0.2)
+        self.rate = jnp.array([[prey_rates[i], predator_rates[i]] for i in range(n_dim)]).ravel()
+        interaction = jnp.zeros((n_var, n_var))
+        for i in range(n_var):
+            a1_key, a2_key = jrandom.split(keys[2+i])
+            a1 = jrandom.uniform(a1_key, minval=0.01, maxval=0.1) * (-1 + 2*((i%2) > 0))
+            interaction = interaction.at[i, (i-1)%n_var].set(a1)
+            a2 = jrandom.uniform(a2_key, minval=0.0, maxval=0.1-a1) * (-1 + 2*((i%2) > 0))
+            interaction = interaction.at[i, (i+1)%n_var].set(a2)
+
+        self.interaction = interaction
+
+        print(self.rate)
+        print(self.interaction)
+        
+        self.V = self.process_noise * jnp.eye(self.n_var)
+        self.W = self.obs_noise * jnp.eye(self.n_obs)[:self.n_obs]
+        self.C = jnp.eye(self.n_var)[:self.n_obs]
+
+    def sample_init_states(self, batch_size, key):
+        return jrandom.uniform(key, shape = (batch_size,self.n_var), minval=5, maxval=15)
+    
+    def sample_init_state2(self, ys, batch_size, key):
+        return ys[jrandom.choice(key, jnp.arange(ys.shape[0]), shape=(batch_size,), replace=False)]
+    
+    def drift(self, t, state, args):
+        # return jnp.array([self.alpha * state[0] - self.beta * state[0] * state[1], self.delta * state[0] * state[1] - self.gamma * state[1]])
+        return state * (self.rate + self.interaction@state)
+
+    def diffusion(self, t, state, args):
+        return self.V
+
+    def terminate_event(self, state, **kwargs):
+        return False

kozax-0.0.1/Kozax/environments/SR_environments/time_series_environment_base.py

@@ -0,0 +1,67 @@
+"""
+Kozax: Genetic programming framework in JAX
+
+Copyright (c) 2024 sdevries0
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with this program.  If not, see <https://www.gnu.org/licenses/>.
+"""
+
+import jax
+import jax.numpy as jnp
+import jax.random as jrandom
+import abc
+
+_itemsize_kind_type = {
+    (1, "i"): jnp.int8,
+    (2, "i"): jnp.int16,
+    (4, "i"): jnp.int32,
+    (8, "i"): jnp.int64,
+    (2, "f"): jnp.float16,
+    (4, "f"): jnp.float32,
+    (8, "f"): jnp.float64,
+}
+
+def force_bitcast_convert_type(val, new_type=jnp.int32):
+    val = jnp.asarray(val)
+    intermediate_type = _itemsize_kind_type[new_type.dtype.itemsize, val.dtype.kind]
+    val = val.astype(intermediate_type)
+    return jax.lax.bitcast_convert_type(val, new_type)
+
+class EnvironmentBase(abc.ABC):
+    def __init__(self, process_noise, obs_noise, n_var, n_obs):
+        self.process_noise = process_noise
+        self.obs_noise = obs_noise
+        self.n_var = n_var
+        self.n_obs = n_obs
+
+    @abc.abstractmethod
+    def sample_init_states(self, batch_size, key):
+        return
+
+    def f_obs(self, key, t_x):
+        t, x = t_x
+        new_key = jrandom.fold_in(key, force_bitcast_convert_type(t))
+        out = self.C@x + jrandom.normal(new_key, shape=(self.n_obs,))@self.W
+        return key, out
+    
+    @abc.abstractmethod
+    def drift(self, t, state, args):
+        return
+
+    @abc.abstractmethod
+    def diffusion(self, t, state, args):
+        return
+
+    def terminate_event(self, state, **kwargs):
+        return False

kozax-0.0.1/Kozax/environments/SR_environments/vd_pol_oscillator.py

@@ -0,0 +1,48 @@
+"""
+Kozax: Genetic programming framework in JAX
+
+Copyright (c) 2024 sdevries0
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with this program.  If not, see <https://www.gnu.org/licenses/>.
+"""
+
+import jax
+import jax.numpy as jnp
+import jax.random as jrandom
+from Kozax.environments.SR_environments.time_series_environment_base import EnvironmentBase
+
+class VanDerPolOscillator(EnvironmentBase):
+    def __init__(self, process_noise, obs_noise, n_obs=2):
+        n_var = 2
+        super().__init__(process_noise, obs_noise, n_var, n_obs)
+
+        self.init_mu = jnp.array([0,0])
+        self.init_sd = jnp.array([1.0,1.0])
+
+        self.mu = 1
+        self.V = self.process_noise * jnp.eye(self.n_var)
+        self.W = self.obs_noise * jnp.eye(self.n_obs)[:self.n_obs]
+        self.C = jnp.eye(self.n_var)[:self.n_obs]
+
+    def sample_init_states(self, batch_size, key):
+        return self.init_mu + self.init_sd*jrandom.normal(key, shape=(batch_size,2))
+    
+    def drift(self, t, state, args):
+        return jnp.array([state[1], self.mu*(1-state[0]**2)*state[1] - state[0]])
+
+    def diffusion(self, t, state, args):
+        return self.V
+
+    def terminate_event(self, state, **kwargs):
+        return False

kozax-0.0.1/Kozax/environments/control_environments/acrobot.py

@@ -0,0 +1,231 @@
+"""
+Kozax: Genetic programming framework in JAX
+
+Copyright (c) 2024 sdevries0
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with this program.  If not, see <https://www.gnu.org/licenses/>.
+"""
+
+import jax
+import jax.numpy as jnp
+import jax.random as jrandom
+
+from Kozax.environments.control_environments.control_environment_base import EnvironmentBase
+
+class Acrobot(EnvironmentBase):
+    def __init__(self, process_noise, obs_noise, n_obs=4):
+        self.n_var = 4
+        self.n_control = 1
+        self.n_targets = 0
+        self.n_dim = 1
+        self.init_bounds = jnp.array([0.1,0.1,0.1,0.1])
+        super().__init__(process_noise, obs_noise, self.n_var, self.n_control, self.n_dim, n_obs)
+
+        self.R = jnp.array([[0.01]])
+
+    def sample_init_states(self, batch_size, key):
+        init_key, target_key = jrandom.split(key)
+        x0 = jrandom.uniform(init_key, shape=(batch_size, self.n_var), minval= -self.init_bounds, maxval= self.init_bounds)
+        targets = jnp.zeros((batch_size, self.n_targets))
+        return x0, targets
+    
+    def sample_params(self, batch_size, mode, ts, key):
+        l1 = l2 = m1 = m2 = jnp.ones((batch_size))
+
+        return l1, l2, m1, m2
+    
+    def f_obs(self, key, t_x):
+        _, out = super().f_obs(key, t_x)
+        out = jnp.array([(out[0]+jnp.pi)%(2*jnp.pi)-jnp.pi, (out[1]+jnp.pi)%(2*jnp.pi)-jnp.pi, out[2], out[3]])[:self.n_obs]
+        return key, out
+
+    def initialize_parameters(self, params, ts):
+        l1, l2, m1, m2 = params
+        self.l1 = l1  # [m]
+        self.l2 = l2  # [m]
+        self.m1 = m1  #: [kg] mass of link 1
+        self.m2 = m2  #: [kg] mass of link 2
+        self.lc1 = 0.5*self.l1  #: [m] position of the center of mass of link 1
+        self.lc2 = 0.5*self.l2  #: [m] position of the center of mass of link 2
+        self.moi1 = self.moi2 = 1.0
+        self.g = 9.81
+
+        self.G = jnp.array([[0,0,0,0],[0,0,0,0],[0,0,1,0],[0,0,0,1]])
+        self.V = self.process_noise*self.G
+
+        self.C = jnp.eye(self.n_var)[:self.n_obs]
+        self.W = self.obs_noise*jnp.eye(self.n_obs)
+    
+    def drift(self, t, state, args):
+        control = jnp.squeeze(args)
+        control = jnp.clip(control, -1, 1)
+        theta1, theta2, theta1_dot, theta2_dot = state
+
+        d1 = self.m1 * self.lc1**2 + self.m2 * (self.l1**2 + self.lc2**2 + 2 * self.l1 * self.lc2 * jnp.cos(theta2)) + self.moi1 + self.moi2
+        d2 = self.m2 * (self.lc2**2 + self.l1 * self.lc2 * jnp.cos(theta2)) + self.moi2
+
+        phi2 = self.m2 * self.lc2 * self.g * jnp.cos(theta1 + theta2 - jnp.pi/2)
+        phi1 = -self.m2 * self.l1 * self.lc2 * theta2_dot**2 * jnp.sin(theta2) - 2 * self.m2 * self.l1 * self.lc2 * theta1_dot * theta2_dot * jnp.sin(theta1) \
+                    + (self.m1 * self.lc1 + self.m2 * self.l1) * self.g * jnp.cos(theta1 - jnp.pi/2) + phi2
+        
+        theta2_acc = (control + d2/d1 * phi1 - self.m2 * self.l1 * self.lc2 * theta1_dot**2 * jnp.sin(theta2) - phi2) \
+                    / (self.m2 * self.lc2**2 + self.moi2 - d2**2 / d1)
+        theta1_acc = -(d2 * theta2_acc + phi1)/d1
+
+        return jnp.array([
+            theta1_dot,
+            theta2_dot,
+            theta1_acc,
+            theta2_acc
+        ])
+
+    def diffusion(self, t, state, args):
+        return self.V
+
+    def fitness_function(self, state, control, target, ts):
+        reached_threshold = jax.vmap(lambda theta1, theta2: -jnp.cos(theta1) - jnp.cos(theta1 + theta2) > 1.5)(state[:,0], state[:,1])
+        first_success = jnp.argmax(reached_threshold)
+
+        control_cost = jax.vmap(lambda _state, _u: _u@self.R@_u)(state, control)
+        costs = jnp.where((ts/(ts[1]-ts[0]))>first_success, jnp.zeros_like(control_cost), control_cost)
+
+        return first_success + (first_success == 0) * ts.shape[0] + jnp.sum(costs)
+    
+    def cond_fn_nan(self, t, y, args, **kwargs):
+        return jnp.where((jnp.abs(y[2])>(8*jnp.pi)) | (jnp.abs(y[3])>(18*jnp.pi)) | jnp.any(jnp.isnan(y)) | jnp.any(jnp.isinf(y)), -1.0, 1.0)
+
+class Acrobot2(EnvironmentBase):
+    def __init__(self, process_noise, obs_noise, n_obs = None):
+        self.n_var = 4
+        self.n_control = 2
+        self.n_targets = 0
+        self.n_dim = 1
+        self.init_bounds = jnp.array([0.1,0.1,0.1,0.1])
+        self.default_obs = 4
+        super().__init__(process_noise, obs_noise, self.n_var, self.n_control, self.n_dim, n_obs if n_obs else self.default_obs)
+
+        self.R = jnp.array(0.01)*jnp.eye(self.n_control)
+
+    def sample_init_states(self, batch_size, key):
+        init_key, target_key = jrandom.split(key)
+        x0 = jrandom.uniform(init_key, shape=(batch_size, self.n_var), minval= -self.init_bounds, maxval= self.init_bounds)
+        targets = jnp.zeros((batch_size, self.n_targets))
+        return x0, targets
+    
+    def sample_params(self, batch_size, mode, ts, key):
+        l1_key, l2_key, m1_key, m2_key, args_key = jrandom.split(key, 5)
+        if mode == "Constant":
+            l1 = l2 = m1 = m2 = jnp.ones((batch_size))
+        elif mode == "Different":
+            l1 = jrandom.uniform(l1_key, shape=(batch_size,), minval=0.75, maxval=1.25)
+            l2 = jrandom.uniform(l2_key, shape=(batch_size,), minval=0.75, maxval=1.25)
+            m1 = jrandom.uniform(m1_key, shape=(batch_size,), minval=0.75, maxval=1.25)
+            m2 = jrandom.uniform(m2_key, shape=(batch_size,), minval=0.75, maxval=1.25)
+        elif mode == "Switch":
+            switch_times = jrandom.randint(args_key, shape=(batch_size,), minval=int(ts.shape[0]/4), maxval=int(3*ts.shape[0]/4))
+            l1 = jnp.zeros((batch_size, ts.shape[0]))
+            l2 = jnp.zeros((batch_size, ts.shape[0]))
+            m1 = jnp.zeros((batch_size, ts.shape[0]))
+            m2 = jnp.zeros((batch_size, ts.shape[0]))
+            for i in range(batch_size):
+                _key11, _key12 = jrandom.split(jrandom.fold_in(l1_key, i))
+                l1 = l1.at[i,:switch_times[i]].set(jrandom.uniform(_key11, shape=(), minval=0.75, maxval=1.25))
+                l1 = l1.at[i,switch_times[i]:].set(jrandom.uniform(_key12, shape=(), minval=0.75, maxval=1.25))
+
+                _key21, _key22 = jrandom.split(jrandom.fold_in(l2_key, i))
+                l2 = l2.at[i,:switch_times[i]].set(jrandom.uniform(_key21, shape=(), minval=0.75, maxval=1.25))
+                l2 = l2.at[i,switch_times[i]:].set(jrandom.uniform(_key22, shape=(), minval=0.75, maxval=1.25))
+
+                _key31, _key32 = jrandom.split(jrandom.fold_in(m1_key, i))
+                m1 = m1.at[i,:switch_times[i]].set(jrandom.uniform(_key31, shape=(), minval=0.75, maxval=1.25))
+                m1 = m1.at[i,switch_times[i]:].set(jrandom.uniform(_key32, shape=(), minval=0.75, maxval=1.25))
+
+                _key41, _key42 = jrandom.split(jrandom.fold_in(m2_key, i))
+                m2 = m2.at[i,:switch_times[i]].set(jrandom.uniform(_key41, shape=(), minval=0.75, maxval=1.25))
+                m2 = m2.at[i,switch_times[i]:].set(jrandom.uniform(_key42, shape=(), minval=0.75, maxval=1.25))
+
+        elif mode == "Decay":
+            decay_factors = jrandom.uniform(args_key, shape=(batch_size,2), minval=0.98, maxval=1.02)
+            init_l1 = jrandom.uniform(l1_key, shape=(batch_size,), minval=0.75, maxval=1.25)
+            init_l2 = jrandom.uniform(l2_key, shape=(batch_size,), minval=0.75, maxval=1.25)
+            init_m1 = jrandom.uniform(m1_key, shape=(batch_size,), minval=0.75, maxval=1.25)
+            init_m2 = jrandom.uniform(m2_key, shape=(batch_size,), minval=0.75, maxval=1.25)
+
+            l1 = jax.vmap(lambda l, d, t: l*(d**t), in_axes=[0, 0, None])(init_l1, decay_factors[:,0], ts)
+            l2 = jax.vmap(lambda l, d, t: l*(d**t), in_axes=[0, 0, None])(init_l2, decay_factors[:,1], ts)
+            m1 = jax.vmap(lambda m, d, t: m*(d**t), in_axes=[0, 0, None])(init_m1, decay_factors[:,0], ts)
+            m2 = jax.vmap(lambda m, d, t: m*(d**t), in_axes=[0, 0, None])(init_m2, decay_factors[:,1], ts)
+
+        return l1, l2, m1, m2
+    
+    def f_obs(self, key, t_x):
+        _, out = super().f_obs(key, t_x)
+        out = jnp.array([(out[0]+jnp.pi)%(2*jnp.pi)-jnp.pi, (out[1]+jnp.pi)%(2*jnp.pi)-jnp.pi, out[2], out[3]])[:self.n_obs]
+        return key, out
+
+    def initialize_parameters(self, params, ts):
+        l1, l2, m1, m2 = params
+        self.l1 = l1  # [m]
+        self.l2 = l2  # [m]
+        self.m1 = m1  #: [kg] mass of link 1
+        self.m2 = m2  #: [kg] mass of link 2
+        self.lc1 = 0.5*self.l1  #: [m] position of the center of mass of link 1
+        self.lc2 = 0.5*self.l2  #: [m] position of the center of mass of link 2
+        self.moi1 = self.moi2 = 1.0
+        self.g = 9.81
+
+        self.G = jnp.array([[0,0,0,0],[0,0,0,0],[0,0,1,0],[0,0,0,1]])
+        self.V = self.process_noise*self.G
+
+        self.C = jnp.eye(self.n_var)[:self.n_obs]
+        self.W = self.obs_noise*jnp.eye(self.n_obs)*(jnp.array([1,1,1,1])[:self.n_obs])
+    
+    def drift(self, t, state, args):
+        control = jnp.squeeze(args)
+        control = jnp.clip(control, -1, 1)
+        c1, c2 = control
+        theta1, theta2, theta1_dot, theta2_dot = state
+
+        d1 = self.m1 * self.lc1**2 + self.m2 * (self.l1**2 + self.lc2**2 + 2 * self.l1 * self.lc2 * jnp.cos(theta2)) + self.moi1 + self.moi2
+        d2 = self.m2 * (self.lc2**2 + self.l1 * self.lc2 * jnp.cos(theta2)) + self.moi2
+
+        phi2 = self.m2 * self.lc2 * self.g * jnp.cos(theta1 + theta2 - jnp.pi/2)
+        phi1 = -self.m2 * self.l1 * self.lc2 * theta2_dot**2 * jnp.sin(theta2) - 2 * self.m2 * self.l1 * self.lc2 * theta1_dot * theta2_dot * jnp.sin(theta1) \
+                    + (self.m1 * self.lc1 + self.m2 * self.l1) * self.g * jnp.cos(theta1 - jnp.pi/2) + phi2
+        
+        theta2_acc = (c1 + d2/d1 * phi1 - self.m2 * self.l1 * self.lc2 * theta1_dot**2 * jnp.sin(theta2) - phi2) \
+                    / (self.m2 * self.lc2**2 + self.moi2 - d2**2 / d1)
+        theta1_acc = (c2 - d2 * theta2_acc - phi1)/d1
+
+        return jnp.array([
+            theta1_dot,
+            theta2_dot,
+            theta1_acc,
+            theta2_acc
+        ])
+
+    def diffusion(self, t, state, args):
+        return self.V
+
+    def fitness_function(self, state, control, target, ts):
+        reached_threshold = jax.vmap(lambda theta1, theta2: -jnp.cos(theta1) - jnp.cos(theta1 + theta2) > 1.5)(state[:,0], state[:,1])
+        first_success = jnp.argmax(reached_threshold)
+
+        control_cost = jax.vmap(lambda _state, _u: _u@self.R@_u)(state, control)
+        costs = jnp.where((ts/(ts[1]-ts[0]))>first_success, jnp.zeros_like(control_cost), control_cost)
+
+        return first_success + (first_success == 0) * ts.shape[0] + jnp.sum(costs)
+
+    def terminate_event(self, state, **kwargs):
+        return (jnp.abs(state.y[2])>(8*jnp.pi)) | (jnp.abs(state.y[3])>(18*jnp.pi)) | jnp.any(jnp.isnan(state.y)) | jnp.any(jnp.isinf(state.y))

kozax-0.0.1/Kozax/environments/control_environments/cart_pole.py

@@ -0,0 +1,99 @@
+"""
+Kozax: Genetic programming framework in JAX
+
+Copyright (c) 2024 sdevries0
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with this program.  If not, see <https://www.gnu.org/licenses/>.
+"""
+
+import jax
+import jax.numpy as jnp
+import jax.random as jrandom
+
+from Kozax.environments.control_environments.control_environment_base import EnvironmentBase
+
+class CartPole(EnvironmentBase):
+    def __init__(self, process_noise, obs_noise, n_obs = 4):
+        self.n_var = 4
+        self.n_control = 1
+        self.n_targets = 0
+        self.n_dim = 1
+        self.init_bounds = jnp.array([0.05,0.05,0.05,0.05])
+        super().__init__(process_noise, obs_noise, self.n_var, self.n_control, self.n_dim, n_obs)
+
+        self.Q = jnp.array(0)
+        self.R = jnp.array([[0.0]])
+
+    def sample_init_states(self, batch_size, key):
+        init_key, target_key = jrandom.split(key)
+        x0 = jrandom.uniform(init_key, shape=(batch_size, self.n_var), minval= -self.init_bounds, maxval= self.init_bounds)
+        targets = jnp.zeros((batch_size, self.n_targets))
+        return x0, targets
+    
+    def sample_params(self, batch_size, mode, ts, key):
+        #g, mass, length
+        return jnp.zeros((batch_size))
+
+    def initialize_parameters(self, params, ts):
+        _ = params
+        self.g = 9.81
+        self.pole_mass = 0.1
+        self.pole_length = 0.5
+        self.cart_mass = 1
+        
+        self.G = jnp.array([[0,0,0,0],[0,0,0,0],[0,0,1,0],[0,0,0,0]])
+        self.V = self.process_noise*self.G
+
+        self.C = jnp.eye(self.n_var)[:self.n_obs]
+        self.W = self.obs_noise*jnp.eye(self.n_obs)
+
+    def drift(self, t, state, args):
+        control = jnp.squeeze(args)
+        control = jnp.clip(control, -1, 1)
+        x, theta, x_dot, theta_dot = state
+
+        cos_theta = jnp.cos(theta)
+        sin_theta = jnp.sin(theta)
+        
+        theta_acc = (self.g * sin_theta - cos_theta * (
+            control + self.pole_mass * self.pole_length * theta_dot**2 * sin_theta
+            ) / (self.cart_mass + self.pole_mass)) / (
+                self.pole_length * (4/3 - (self.pole_mass * cos_theta**2) / (self.cart_mass + self.pole_mass)))
+
+        x_acc = (control + self.pole_mass * self.pole_length * (theta_dot**2 * sin_theta - theta_acc * cos_theta)) / (self.cart_mass + self.pole_mass)
+
+        # x_acc = control
+        # theta_acc = (self.g/self.pole_length) * jnp.sin(theta)
+
+        return jnp.array([
+            x_dot,
+            theta_dot,
+            x_acc,
+            theta_acc
+        ])
+    
+    def diffusion(self, t, x, args):
+        return self.V
+    
+    def fitness_function(self, state, control, target, ts):
+        invalid_points = jax.vmap(lambda _x, _u: jnp.any(jnp.isinf(_x)) + jnp.isnan(_u))(state, control[:,0])
+        punishment = jnp.ones_like(invalid_points)
+        # costs_control = jax.vmap(lambda _state, _u: _u@self.R@_u)(state, u)
+
+        costs = jnp.where(invalid_points, punishment, jnp.zeros_like(punishment))
+
+        return jnp.sum(costs)
+    
+    def terminate_event(self, state, **kwargs):
+        return (jnp.abs(state.y[1])>0.2) | (jnp.abs(state.y[0])>4.8) | jnp.any(jnp.isnan(state.y)) | jnp.any(jnp.isinf(state.y))