spiderbot-simulation/world_simulate.py at master · matagno/spiderbot-simulation · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
import pybullet as p
import pybullet_data
import numpy as np

from robot import Robot
from interface import Interface

class WorldSimulate:

    def __init__(self, client_type, env_path, robot_path, heightfield=False):
        self.client_type = client_type
        self.client_id = p.connect(client_type)

        if heightfield :
            self.ground_id = self.create_heightfield_ground()
        else :
            self.ground_id = self.load_world(env_path)
        self.bot = self.load_object(robot_path)

        self.interface = Interface(self.client_id) if client_type == p.GUI else None


    def load_world(self, env_path):
        p.setAdditionalSearchPath(pybullet_data.getDataPath(), physicsClientId=self.client_id)
        p.setGravity(0, 0, -9.81, physicsClientId=self.client_id)

        ground_id = p.loadURDF(env_path, physicsClientId=self.client_id)

        return ground_id


    def create_heightfield_ground(self):
        p.setGravity(0, 0, -9.81)

        # Dimensions
        numRows = 64
        numCols = 64
        heightfieldData = np.zeros((numRows, numCols))

        # Zone
        left_cols = numCols // 4
        middle_cols = numCols // 2
        right_cols = numCols - (left_cols + middle_cols)
        bump_amplitude = 0.01

        # Zone left : Up
        for j in range(left_cols):
            heightfieldData[:, j] = 0.01 * j + np.random.uniform(-bump_amplitude, bump_amplitude, size=numRows)

        # Zone center : flat
        flat_height = 0.01 * (left_cols - 1)
        heightfieldData[:, left_cols:left_cols+middle_cols] = flat_height

        # Zone right : Up
        for j in range(right_cols):
            start_height = heightfieldData[:, left_cols + middle_cols - 1]
            heightfieldData[:, left_cols + middle_cols + j] = start_height + 0.01 * j + np.random.uniform(-bump_amplitude, bump_amplitude, size=numRows)

        # Field
        field_shape = p.createCollisionShape(
            shapeType=p.GEOM_HEIGHTFIELD,
            meshScale=[0.05, 0.05, 1],
            heightfieldTextureScaling=(numRows - 1)/2,
            heightfieldData=heightfieldData.flatten(),
            numHeightfieldRows=numRows,
            numHeightfieldColumns=numCols
        )

        field = p.createMultiBody(0, field_shape, basePosition=[0,0,0])
        return field


    def load_object(self, robot_path):
        robot_start_pos = [0, 0, 0.8]  # Augmentation de la hauteur initiale
        robot_start_orientation = p.getQuaternionFromEuler([0, 0, 0])

        bot = Robot(self.client_id,
                        robot_path,
                        robot_start_pos,
                        robot_start_orientation)

        p.changeVisualShape(bot.id, -1, rgbaColor=[0.6, 0, 0, 1], physicsClientId=self.client_id)
        for i in range(0, 16) :
            if i%2 == 0:
                p.changeVisualShape(bot.id, i, rgbaColor=[0, 0, 0, 1], physicsClientId=self.client_id)  # Noir
            else :
                p.changeVisualShape(bot.id, i, rgbaColor=[0.6, 0, 0, 1], physicsClientId=self.client_id)  # Rouge


        return bot