test_PolicyIterationModified.py 755 Bytes
Newer Older
Steven Cordwell's avatar
Steven Cordwell committed
1 2 3 4 5 6 7
# -*- coding: utf-8 -*-
"""
Created on Wed Sep 11 13:33:16 2013

@author: steve
"""

8 9 10 11 12 13
from nose.tools import assert_equal
import numpy as np
import scipy.sparse as sp

from mdptoolbox import mdp

14
from .utils import BaseTestMDP, assert_sequence_almost_equal
15 16 17 18 19 20 21 22 23 24 25 26 27

class TestPolicyIterationModified(BaseTestMDP):
    def test_small(self):
        pim = mdp.PolicyIterationModified(self.small_P, self.small_R, 0.9)
        pim.run()
        assert_sequence_almost_equal(pim.V,
                                    (41.8656419239403, 35.4702797722819))
        assert_equal(pim.policy, (1, 0))

    def test_small_undiscounted(self):
        pim = mdp.PolicyIterationModified(self.small_P, self.small_R, 1)
        pim.run()
        assert_equal(pim.policy, (1, 0))