Skip to content

Commit

Permalink
Fix/set mean reversion target to D0(4)
Browse files Browse the repository at this point in the history
Supersedes and closes #123
  • Loading branch information
user1823 authored Jul 26, 2024
1 parent 86eeb69 commit 9a3494c
Showing 1 changed file with 6 additions and 2 deletions.
8 changes: 6 additions & 2 deletions src/fsrs_optimizer/fsrs_optimizer.py
Original file line number Diff line number Diff line change
Expand Up @@ -95,9 +95,13 @@ def stability_short_term(self, state: Tensor, rating: Tensor) -> Tensor:
new_s = state[:, 0] * torch.exp(self.w[17] * (rating - 3 + self.w[18]))
return new_s

def init_d(self, rating: Tensor) -> Tensor:
new_d = self.w[4] - torch.exp(self.w[5] * (X[:, 1] - 1)) + 1
return new_d

def next_d(self, state: Tensor, rating: Tensor) -> Tensor:
new_d = state[:, 1] - self.w[6] * (rating - 3)
new_d = self.mean_reversion(self.w[4], new_d)
new_d = self.mean_reversion(init_d(4), new_d)
return new_d

def step(self, X: Tensor, state: Tensor) -> Tensor:
Expand All @@ -113,7 +117,7 @@ def step(self, X: Tensor, state: Tensor) -> Tensor:
# first learn, init memory states
new_s = torch.ones_like(state[:, 0])
new_s[index[0]] = self.w[index[1]]
new_d = self.w[4] - torch.exp(self.w[5] * (X[:, 1] - 1)) + 1
new_d = init_d(X[:, 1])
new_d = new_d.clamp(1, 10)
else:
r = power_forgetting_curve(X[:, 0], state[:, 0])
Expand Down

0 comments on commit 9a3494c

Please sign in to comment.